4th commit
Refactored some code