Bugfix and minor code cleanup.