New code for lab 6