Clean up code base on code review