Fix regressions due to code cleanup