cleanup the code

(forgot to amend previous commit)