Refactored logic