More code cleanup (not finished)