More code cleaning up