Cleaned up the code generator