Refactor code iteration 1