Minor refactoring and style update