code refactor to make UML look more pretty