slight class diagram changes