Minor refactoring and tweaks