Code cleanup