Added support for return statements
Code formating