More code cleanup