Code cleanup refactoring