Fix post-refactoring errors