Added the square function. Woo I love life