Refactoring and some performace improvements