Refactor the code by moving things around

This just tries to get rid of `case` soup.