Fix bugs introduced by refactor.