Some more code cleanup