Minor tweaks and refactoring