Code cleanup again