Some code cleanup and review