Fix missing refactored logic from previous commit.