Refactor code

Here are the refactorings I was talking about.