Updated diagram
Some more code reorganizations