Controller code clean up