redesign according to class diagram