Additional cleanup of diagram tool