Minor fixes and refactors