Code clean up and small optimizations