Did some cleanup and minor changes after refactoring.