Add refactor suggestions from Puneeth