Minor code cleanup

I simplified a few if statements and ran `go fmt` on the files that I modified. 
(I realise you just copied it across and didn't write it yourself)