Fixed issue #22
Minor changes to code generation