More refactor work