Code cleanup and structure