Code cleanup
Bit more cleanup