Bug fixes + a bit of code cleanup.