refactor code for new module code