class diagram fix