Added some null ref checks and code clean up.