Clean up and small bug fix.