Add some comments and small code cleanup.