Code cleanup

Just some code cleanup.