Slight refactoring and clean up.