Handle nested sequence diagrams

d2layouts/d2sequence/layout.go

@@ -5,44 +5,109 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"strings"
 
 	"oss.terrastruct.com/d2/d2graph"
+	"oss.terrastruct.com/d2/d2target"
 	"oss.terrastruct.com/d2/lib/geo"
 	"oss.terrastruct.com/d2/lib/go2"
 	"oss.terrastruct.com/d2/lib/label"
 	"oss.terrastruct.com/d2/lib/shape"
 )
 
-func Layout2(ctx context.Context, g *d2graph.Graph, layoutFn func(ctx context.Context, g *d2graph.Graph) error) error {
+func Layout2(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Context, g *d2graph.Graph) error) error {
+	// new graph objects without sequence diagram objects and their replacement (rectangle node)
+	newObjects := make([]*d2graph.Object, 0, len(g.Objects))
+	edgesToRemove := make(map[*d2graph.Edge]struct{})
+
+	sequenceDiagrams := make(map[*d2graph.Object]*sequenceDiagram)
+
+	objChildrenArray := make(map[*d2graph.Object][]*d2graph.Object)
+
+	queue := make([]*d2graph.Object, 1, len(g.Objects))
+	queue[0] = g.Root
+	for len(queue) > 0 {
+		obj := queue[0]
+		queue = queue[1:]
+
+		newObjects = append(newObjects, obj)
+		if obj.Attributes.Shape.Value == d2target.ShapeSequenceDiagram {
+			// TODO: should update obj.References too?
+
+			// clean current children and keep a backup to restore them later
+			obj.Children = make(map[string]*d2graph.Object)
+			objChildrenArray[obj] = obj.ChildrenArray
+			obj.ChildrenArray = nil
+			// creates a mock rectangle so that layout considers this size
+			sdMock := obj.EnsureChild([]string{"sequence_diagram"})
+			sdMock.Attributes.Shape.Value = d2target.ShapeRectangle
+			sdMock.Attributes.Label.Value = ""
+			newObjects = append(newObjects, sdMock)
+
+			var messages []*d2graph.Edge
+			for _, edge := range g.Edges {
+				if strings.HasPrefix(edge.Src.AbsID(), obj.AbsID()) && strings.HasPrefix(edge.Dst.AbsID(), obj.AbsID()) {
+					edgesToRemove[edge] = struct{}{}
+					messages = append(messages, edge)
+				}
+			}
+
+			sd := newSequenceDiagram(objChildrenArray[obj], messages)
+			sd.layout()
+			sdMock.Box = geo.NewBox(nil, sd.getWidth(), sd.getHeight())
+			sequenceDiagrams[obj] = sd
+		} else {
+			queue = append(queue, obj.ChildrenArray...)
+		}
+	}
+
+	newEdges := make([]*d2graph.Edge, 0, len(g.Edges)-len(edgesToRemove))
+	for _, edge := range g.Edges {
+		if _, exists := edgesToRemove[edge]; !exists {
+			newEdges = append(newEdges, edge)
+		}
+	}
+
+	g.Objects = newObjects
+	g.Edges = newEdges
+
+	if err := layout(ctx, g); err != nil {
+		return err
+	}
+
+	// restores objects & edges
+	for edge := range edgesToRemove {
+		g.Edges = append(g.Edges, edge)
+	}
+	for obj, children := range objChildrenArray {
+		sdMock := obj.ChildrenArray[0]
+		sequenceDiagrams[obj].shift(sdMock.TopLeft)
+		obj.Children = make(map[string]*d2graph.Object)
+		for _, child := range children {
+			g.Objects = append(g.Objects, child)
+			obj.Children[child.ID] = child
+		}
+		obj.ChildrenArray = children
+
+		for _, edge := range sequenceDiagrams[obj].lifelines {
+			g.Edges = append(g.Edges, edge)
+		}
+	}
+
 	return nil
 }
 
 func Layout(ctx context.Context, g *d2graph.Graph) (err error) {
-	sd := &sequenceDiagram{
+	sd := newSequenceDiagram(nil, nil)
-		graph:          g,
+	sd.layout()
-		objectRank:     make(map[*d2graph.Object]int),
-		minMessageRank: make(map[*d2graph.Object]int),
-		maxMessageRank: make(map[*d2graph.Object]int),
-		messageYStep:   MIN_MESSAGE_DISTANCE,
-		actorXStep:     MIN_ACTOR_DISTANCE,
-		maxActorHeight: 0.,
-	}
-
-	sd.init()
-	sd.placeActors()
-	sd.placeSpans()
-	sd.routeMessages()
-	sd.addLifelineEdges()
-
 	return nil
 }
 
 type sequenceDiagram struct {
-	graph *d2graph.Graph
+	messages  []*d2graph.Edge
-
+	lifelines []*d2graph.Edge
-	messages []*d2graph.Edge
+	actors    []*d2graph.Object
-	actors   []*d2graph.Object
+	spans     []*d2graph.Object
-	spans    []*d2graph.Object
 
 	// can be either actors or spans
 	// rank: left to right position of actors/spans (spans have the same rank as their parents)
@@ -58,29 +123,38 @@ type sequenceDiagram struct {
 	maxActorHeight float64
 }
 
-func (sd *sequenceDiagram) init() {
+func newSequenceDiagram(actors []*d2graph.Object, messages []*d2graph.Edge) *sequenceDiagram {
-	sd.messages = make([]*d2graph.Edge, len(sd.graph.Edges))
+	sd := &sequenceDiagram{
-	copy(sd.messages, sd.graph.Edges)
+		messages:       messages,
+		actors:         actors,
+		spans:          nil,
+		lifelines:      nil,
+		objectRank:     make(map[*d2graph.Object]int),
+		minMessageRank: make(map[*d2graph.Object]int),
+		maxMessageRank: make(map[*d2graph.Object]int),
+		messageYStep:   MIN_MESSAGE_DISTANCE,
+		actorXStep:     MIN_ACTOR_DISTANCE,
+		maxActorHeight: 0.,
+	}
 
-	queue := make([]*d2graph.Object, len(sd.graph.Root.ChildrenArray))
+	for rank, actor := range actors {
-	copy(queue, sd.graph.Root.ChildrenArray)
+		sd.objectRank[actor] = rank
-	for len(queue) > 0 {
+		sd.maxActorHeight = math.Max(sd.maxActorHeight, actor.Height)
-		obj := queue[0]
+
-		queue = queue[1:]
+		queue := make([]*d2graph.Object, len(actor.ChildrenArray))
+		copy(queue, actor.ChildrenArray)
+		for len(queue) > 0 {
+			span := queue[0]
+			queue = queue[1:]
 
-		if sd.isActor(obj) {
-			sd.actors = append(sd.actors, obj)
-			sd.objectRank[obj] = len(sd.actors)
-			sd.maxActorHeight = math.Max(sd.maxActorHeight, obj.Height)
-		} else {
 			// spans are always rectangles and have no labels
-			obj.Attributes.Label = d2graph.Scalar{Value: ""}
+			span.Attributes.Label = d2graph.Scalar{Value: ""}
-			obj.Attributes.Shape = d2graph.Scalar{Value: shape.SQUARE_TYPE}
+			span.Attributes.Shape = d2graph.Scalar{Value: shape.SQUARE_TYPE}
-			sd.spans = append(sd.spans, obj)
+			sd.spans = append(sd.spans, span)
-			sd.objectRank[obj] = sd.objectRank[obj.Parent]
+			sd.objectRank[span] = rank
-		}
 
-		queue = append(queue, obj.ChildrenArray...)
+			queue = append(queue, span.ChildrenArray...)
+		}
 	}
 
 	for rank, message := range sd.messages {
@@ -98,6 +172,8 @@ func (sd *sequenceDiagram) init() {
 
 	sd.maxActorHeight += VERTICAL_PAD
 	sd.messageYStep += VERTICAL_PAD
+
+	return sd
 }
 
 func (sd *sequenceDiagram) setMinMaxMessageRank(actor *d2graph.Object, rank int) {
@@ -110,6 +186,13 @@ func (sd *sequenceDiagram) setMinMaxMessageRank(actor *d2graph.Object, rank int)
 	sd.maxMessageRank[actor] = go2.IntMax(sd.maxMessageRank[actor], rank)
 }
 
+func (sd *sequenceDiagram) layout() {
+	sd.placeActors()
+	sd.placeSpans()
+	sd.routeMessages()
+	sd.addLifelineEdges()
+}
+
 // placeActors places actors bottom aligned, side by side
 func (sd *sequenceDiagram) placeActors() {
 	x := 0.
@@ -136,7 +219,7 @@ func (sd *sequenceDiagram) addLifelineEdges() {
 		actorBottom.Y = actor.TopLeft.Y + actor.Height
 		actorLifelineEnd := actor.Center()
 		actorLifelineEnd.Y = endY
-		sd.graph.Edges = append(sd.graph.Edges, &d2graph.Edge{
+		sd.lifelines = append(sd.lifelines, &d2graph.Edge{
 			Attributes: d2graph.Attributes{
 				Style: d2graph.Style{
 					StrokeDash:  &d2graph.Scalar{Value: "10"},
@@ -280,5 +363,40 @@ func (sd *sequenceDiagram) getMessageY(rank int) float64 {
 }
 
 func (sd *sequenceDiagram) isActor(obj *d2graph.Object) bool {
-	return obj.Parent == sd.graph.Root
+	// TODO: map to avoid looping around every time?
+	for _, actor := range sd.actors {
+		if actor == obj {
+			return true
+		}
+	}
+	return false
+}
+
+func (sd *sequenceDiagram) getWidth() float64 {
+	// the layout is always placed starting at 0, so the width is just the last actor
+	lastActor := sd.actors[len(sd.actors)-1]
+	return lastActor.TopLeft.X + lastActor.Width
+}
+
+func (sd *sequenceDiagram) getHeight() float64 {
+	// the layout is always placed starting at 0, so the height is just the last message
+	return sd.getMessageY(len(sd.messages))
+}
+
+func (sd *sequenceDiagram) shift(tl *geo.Point) {
+	allObjects := append([]*d2graph.Object{}, sd.actors...)
+	allObjects = append(allObjects, sd.spans...)
+	for _, obj := range allObjects {
+		obj.TopLeft.X += tl.X
+		obj.TopLeft.Y += tl.Y
+	}
+
+	allEdges := append([]*d2graph.Edge{}, sd.messages...)
+	allEdges = append(allEdges, sd.lifelines...)
+	for _, edge := range allEdges {
+		for _, p := range edge.Route {
+			p.X += tl.X
+			p.Y += tl.Y
+		}
+	}
 }

d2layouts/d2sequence/layout_test.go

@@ -257,15 +257,18 @@ func TestSpansSequenceDiagram(t *testing.T) {
 }
 
 func TestNestedSequenceDiagrams(t *testing.T) {
-	//   ┌─────┐                 ┌─────┐
+	// ┌────────────────────────────────────────┐
-	//   │  a  │                 │  b  │
+	// |     ┌─────┐    container    ┌─────┐    |
-	//   └──┬──┘                 └──┬──┘
+	// |     │  a  │                 │  b  │    |            ┌─────┐
-	//      ├┐────────────────────►┌┤
+	// |     └──┬──┘                 └──┬──┘    ├────edge1───┤  c  │
-	//   t1 ││                     ││ t1
+	// |        ├┐───────sdEdge1──────►┌┤       |            └─────┘
-	//      ├┘◄────────────────────└┤
+	// |     t1 ││                     ││ t1    |
+	// |        ├┘◄──────sdEdge2───────└┤       |
+	// └────────────────────────────────────────┘
 	g := d2graph.NewGraph(nil)
 	container := g.Root.EnsureChild([]string{"container"})
 	container.Attributes.Shape = d2graph.Scalar{Value: d2target.ShapeSequenceDiagram}
+	container.Box = geo.NewBox(nil, 500, 500)
 	a := container.EnsureChild([]string{"a"})
 	a.Box = geo.NewBox(nil, 100, 100)
 	a.Attributes.Shape = d2graph.Scalar{Value: shape.PERSON_TYPE}
@@ -275,6 +278,7 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 	b_t1 := b.EnsureChild([]string{"t1"})
 
 	c := g.Root.EnsureChild([]string{"c"})
+	c.Box = geo.NewBox(nil, 100, 100)
 	c.Attributes.Shape = d2graph.Scalar{Value: d2target.ShapeSquare}
 
 	sdEdge1, err := g.Root.Connect(a_t1.AbsIDArray(), b_t1.AbsIDArray(), false, true, "sequence diagram edge 1")
@@ -292,6 +296,23 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 	}
 
 	layoutFn := func(ctx context.Context, g *d2graph.Graph) error {
+		// 4 because it replaces all `container` children with a rectangle for layout
+		if len(g.Objects) != 4 {
+			t.Fatal("expected only diagram objects for layout")
+		}
+		for _, obj := range g.Objects {
+			if obj == a || obj == a_t1 || obj == b || obj == b_t1 {
+				t.Fatal("expected to have removed all sequence diagram objects")
+			}
+		}
+		if len(container.ChildrenArray) != 1 {
+			t.Fatalf("expected only 1 `container` child, got %d", len(container.ChildrenArray))
+		}
+
+		if len(container.Children) != len(container.ChildrenArray) {
+			t.Fatal("container children mismatch")
+		}
+
 		for _, edge := range g.Edges {
 			if edge == sdEdge1 || edge == sdEdge2 {
 				t.Fatal("expected to have removed all sequence diagram edges from graph")
@@ -300,6 +321,18 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 		if g.Edges[0] != edge1 {
 			t.Fatal("expected graph edge to be in the graph")
 		}
+
+		// just set some position as if it had been properly placed
+		for _, obj := range g.Objects {
+			if obj != g.Root {
+				obj.TopLeft = geo.NewPoint(0, 0)
+			}
+		}
+
+		for _, edge := range g.Edges {
+			edge.Route = []*geo.Point{geo.NewPoint(1, 1)}
+		}
+
 		return nil
 	}
 
@@ -308,7 +341,18 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if len(g.Edges) != 3 {
+	if len(g.Edges) != 5 {
-		t.Fatal("expected graph to have all edges after layout")
+		t.Fatal("expected graph to have all edges and lifelines after layout")
+	}
+
+	for _, obj := range g.Objects {
+		if obj != g.Root && obj.TopLeft == nil {
+			t.Fatal("expected to have placed all objects")
+		}
+	}
+	for _, edge := range g.Edges {
+		if len(edge.Route) == 0 {
+			t.Fatal("expected to have routed all edges")
+		}
 	}
 }