Minor refactoring and code clean up.