A bit more refactoring... there were still some errors after the last commit.