Update to code generator