Minor changes to code organization