Fixed my bad commit from a while back

Previously when fixing something my files got badly merged, then they
were committed. So I went back and got the correct files from Will's
commits to fix these.