Code cleanup and slight bugfix since last commit (derp moment)