Code made better
one more refactoring