Code cleanup and optimization