Revise assignment 2