Refactor (step 2)