Minor bug fix and code cleanup.