Some minor refactoring and re-ordering