More code clean-up from review