Code clean up pass