Refactor the control flow