Fixing bugs introduced during code cleanup