Fixes for new module codegen