Some minor refactoring / cleanup