Some minor refactoring and cleanup