Fixing and cleaning up some code.