Code cleanup
Fixed bug introduced with code cleanups
Merge branch 'cleanup'