Small cleanup and bug fix.