Cleanup
Code refactoring