Fix some regressions after refactoring