Re-factor previous commit to make code tidier
Minor hlint suggestions