Fixed a few build errors after refactoring