Slight code cleanup, made a switch-case a bit clearer