Updating the controller to simplify the code.