Added code in place of TODO for the lab2 assignment