Refactor return statement checks.