More refactoring fixes