Refactored some code in controller