Add lifespan

d2layouts/d2sequence/constants.go

@@ -7,3 +7,5 @@ const MIN_ACTOR_DISTANCE = 200.
 
 // min vertical distance between edges
 const MIN_EDGE_DISTANCE = 100.
+
+const LIFESPAN_BOX_WIDTH = 20.

d2layouts/d2sequence/layout.go

@@ -14,29 +14,20 @@ import (
 func Layout(ctx context.Context, g *d2graph.Graph) (err error) {
 	sd := &sequenceDiagram{
 		graph:          g,
+		objectRank:     make(map[*d2graph.Object]int),
+		edgeRank:       make(map[*d2graph.Edge]int),
+		minEdgeRank:    make(map[*d2graph.Object]int),
+		maxEdgeRank:    make(map[*d2graph.Object]int),
 		edgeYStep:      MIN_EDGE_DISTANCE,
 		actorXStep:     MIN_ACTOR_DISTANCE,
 		maxActorHeight: 0.,
 	}
 
-	actorRank := make(map[*d2graph.Object]int)
+	sd.init()
-	for rank, actor := range g.Objects {
-		actorRank[actor] = rank
-	}
-	for _, edge := range g.Edges {
-		sd.edgeYStep = math.Max(sd.edgeYStep, float64(edge.LabelDimensions.Height)+HORIZONTAL_PAD)
-		sd.maxActorHeight = math.Max(sd.maxActorHeight, edge.Src.Height+HORIZONTAL_PAD)
-		sd.maxActorHeight = math.Max(sd.maxActorHeight, edge.Dst.Height+HORIZONTAL_PAD)
-		// ensures that long labels, spanning over multiple actors, don't make for large gaps between actors
-		// by distributing the label length across the actors rank difference
-		rankDiff := math.Abs(float64(actorRank[edge.Src]) - float64(actorRank[edge.Dst]))
-		distributedLabelWidth := float64(edge.LabelDimensions.Width) / rankDiff
-		sd.actorXStep = math.Max(sd.actorXStep, distributedLabelWidth+HORIZONTAL_PAD)
-	}
-
 	sd.placeActors()
-	sd.routeEdges()
 	sd.addLifelineEdges()
+	sd.placeLifespan()
+	sd.routeEdges()
 
 	return nil
 }
@@ -44,15 +35,88 @@ func Layout(ctx context.Context, g *d2graph.Graph) (err error) {
 type sequenceDiagram struct {
 	graph *d2graph.Graph
 
+	edges     []*d2graph.Edge
+	actors    []*d2graph.Object
+	lifespans []*d2graph.Object
+
+	// can be either actors or lifespans
+	objectRank map[*d2graph.Object]int
+	edgeRank   map[*d2graph.Edge]int
+
+	// keep track of the first and last edge of a given actor
+	// needed for lifespan
+	minEdgeRank map[*d2graph.Object]int
+	maxEdgeRank map[*d2graph.Object]int
+
 	edgeYStep      float64
 	actorXStep     float64
 	maxActorHeight float64
 }
 
+func intMin(a, b int) int {
+	return int(math.Min(float64(a), float64(b)))
+}
+
+func intMax(a, b int) int {
+	return int(math.Max(float64(a), float64(b)))
+}
+
+func (sd *sequenceDiagram) init() {
+	sd.edges = make([]*d2graph.Edge, len(sd.graph.Edges))
+	copy(sd.edges, sd.graph.Edges)
+
+	for rank, actor := range sd.graph.Root.ChildrenArray {
+		sd.assignRank(actor, rank)
+	}
+	for _, obj := range sd.graph.Objects {
+		if obj.Parent == sd.graph.Root {
+			sd.actors = append(sd.actors, obj)
+		} else if obj != sd.graph.Root {
+			sd.lifespans = append(sd.lifespans, obj)
+		}
+	}
+	for rank, edge := range sd.edges {
+		sd.edgeRank[edge] = rank
+		if edge.Src.Parent == sd.graph.Root {
+			sd.maxActorHeight = math.Max(sd.maxActorHeight, edge.Src.Height+HORIZONTAL_PAD)
+		}
+		if edge.Dst.Parent == sd.graph.Root {
+			sd.maxActorHeight = math.Max(sd.maxActorHeight, edge.Dst.Height+HORIZONTAL_PAD)
+		}
+		sd.edgeYStep = math.Max(sd.edgeYStep, float64(edge.LabelDimensions.Height)+HORIZONTAL_PAD)
+
+		sd.setMinMaxEdgeRank(edge.Src, rank)
+		sd.setMinMaxEdgeRank(edge.Dst, rank)
+
+		// ensures that long labels, spanning over multiple actors, don't make for large gaps between actors
+		// by distributing the label length across the actors rank difference
+		rankDiff := math.Abs(float64(sd.objectRank[edge.Src]) - float64(sd.objectRank[edge.Dst]))
+		distributedLabelWidth := float64(edge.LabelDimensions.Width) / rankDiff
+		sd.actorXStep = math.Max(sd.actorXStep, distributedLabelWidth+HORIZONTAL_PAD)
+	}
+}
+
+func (sd *sequenceDiagram) assignRank(actor *d2graph.Object, rank int) {
+	sd.objectRank[actor] = rank
+	for _, child := range actor.Children {
+		sd.assignRank(child, rank)
+	}
+}
+
+func (sd *sequenceDiagram) setMinMaxEdgeRank(actor *d2graph.Object, rank int) {
+	if minRank, exists := sd.minEdgeRank[actor]; exists {
+		sd.minEdgeRank[actor] = intMin(minRank, rank)
+	} else {
+		sd.minEdgeRank[actor] = rank
+	}
+
+	sd.maxEdgeRank[actor] = intMax(sd.maxEdgeRank[actor], rank)
+}
+
 // placeActors places actors bottom aligned, side by side
 func (sd *sequenceDiagram) placeActors() {
 	x := 0.
-	for _, actors := range sd.graph.Objects {
+	for _, actors := range sd.actors {
 		yOffset := sd.maxActorHeight - actors.Height
 		actors.TopLeft = geo.NewPoint(x, yOffset)
 		x += actors.Width + sd.actorXStep
@@ -60,28 +124,6 @@ func (sd *sequenceDiagram) placeActors() {
 	}
 }
 
-// routeEdges routes horizontal edges from Src to Dst
-func (sd *sequenceDiagram) routeEdges() {
-	edgeY := sd.maxActorHeight + sd.edgeYStep // in case the first edge has a tall label
-	for _, edge := range sd.graph.Edges {
-		start := edge.Src.Center()
-		start.Y = edgeY
-		end := edge.Dst.Center()
-		end.Y = edgeY
-		edge.Route = []*geo.Point{start, end}
-		edgeY += sd.edgeYStep
-
-		if edge.Attributes.Label.Value != "" {
-			isLeftToRight := edge.Src.TopLeft.X < edge.Dst.TopLeft.X
-			if isLeftToRight {
-				edge.LabelPosition = go2.Pointer(string(label.OutsideTopCenter))
-			} else {
-				edge.LabelPosition = go2.Pointer(string(label.OutsideBottomCenter))
-			}
-		}
-	}
-}
-
 // addLifelineEdges adds a new edge for each actor in the graph that represents the
 // edge below the actor showing its lifespan
 // ┌──────────────┐
@@ -92,8 +134,8 @@ func (sd *sequenceDiagram) routeEdges() {
 //        │
 //        │
 func (sd *sequenceDiagram) addLifelineEdges() {
-	endY := sd.graph.Edges[len(sd.graph.Edges)-1].Route[0].Y + sd.edgeYStep
+	endY := sd.getEdgeY(len(sd.edges))
-	for _, actor := range sd.graph.Objects {
+	for _, actor := range sd.actors {
 		actorBottom := actor.Center()
 		actorBottom.Y = actor.TopLeft.Y + actor.Height
 		actorLifelineEnd := actor.Center()
@@ -116,3 +158,44 @@ func (sd *sequenceDiagram) addLifelineEdges() {
 		})
 	}
 }
+
+func (sd *sequenceDiagram) placeLifespan() {
+	rankToX := make(map[int]float64)
+	for _, actor := range sd.actors {
+		rankToX[sd.objectRank[actor]] = actor.Center().X
+	}
+	for _, lifespan := range sd.lifespans {
+		minRank := sd.minEdgeRank[lifespan]
+		maxRank := sd.maxEdgeRank[lifespan]
+
+		minY := sd.getEdgeY(minRank)
+		maxY := sd.getEdgeY(maxRank)
+		height := maxY - minY
+		x := rankToX[sd.objectRank[lifespan]] - (LIFESPAN_BOX_WIDTH / 2.)
+		lifespan.Box = geo.NewBox(geo.NewPoint(x, minY), LIFESPAN_BOX_WIDTH, height)
+	}
+}
+
+// routeEdges routes horizontal edges from Src to Dst
+func (sd *sequenceDiagram) routeEdges() {
+	for rank, edge := range sd.edges {
+		start := edge.Src.Center()
+		start.Y = sd.getEdgeY(rank)
+		end := edge.Dst.Center()
+		end.Y = start.Y
+		edge.Route = []*geo.Point{start, end}
+
+		if edge.Attributes.Label.Value != "" {
+			isLeftToRight := edge.Src.TopLeft.X < edge.Dst.TopLeft.X
+			if isLeftToRight {
+				edge.LabelPosition = go2.Pointer(string(label.OutsideTopCenter))
+			} else {
+				edge.LabelPosition = go2.Pointer(string(label.OutsideBottomCenter))
+			}
+		}
+	}
+}
+
+func (sd *sequenceDiagram) getEdgeY(rank int) float64 {
+	return ((float64(rank) + 1.) * sd.edgeYStep) + sd.maxActorHeight
+}

d2layouts/d2sequence/layout_test.go

@@ -9,35 +9,41 @@ import (
 	"oss.terrastruct.com/d2/lib/log"
 )
 
-func TestLayout(t *testing.T) {
+func TestBasicSequenceDiagram(t *testing.T) {
+	// ┌────────┐              ┌────────┐
+	// │   n1   │              │   n2   │
+	// └────┬───┘              └────┬───┘
+	//      │                       │
+	//      ├───────────────────────►
+	//      │                       │
+	//      ◄───────────────────────┤
+	//      │                       │
+	//      ├───────────────────────►
+	//      │                       │
+	//      ◄───────────────────────┤
+	//      │                       │
 	g := d2graph.NewGraph(nil)
-	g.Objects = []*d2graph.Object{
+	n1 := g.Root.EnsureChild([]string{"n1"})
-		{
+	n1.Box = geo.NewBox(nil, 100, 100)
-			ID:  "Alice",
+	n2 := g.Root.EnsureChild([]string{"n2"})
-			Box: geo.NewBox(nil, 100, 100),
+	n2.Box = geo.NewBox(nil, 30, 30)
-		},
-		{
-			ID:  "Bob",
-			Box: geo.NewBox(nil, 30, 30),
-		},
-	}
 
 	g.Edges = []*d2graph.Edge{
 		{
-			Src: g.Objects[0],
+			Src: n1,
-			Dst: g.Objects[1],
+			Dst: n2,
 		},
 		{
-			Src: g.Objects[1],
+			Src: n2,
-			Dst: g.Objects[0],
+			Dst: n1,
 		},
 		{
-			Src: g.Objects[0],
+			Src: n1,
-			Dst: g.Objects[1],
+			Dst: n2,
 		},
 		{
-			Src: g.Objects[1],
+			Src: n2,
-			Dst: g.Objects[0],
+			Dst: n1,
 		},
 	}
 	nEdges := len(g.Edges)
@@ -110,3 +116,52 @@ func TestLayout(t *testing.T) {
 		}
 	}
 }
+
+func TestLifespanSequenceDiagram(t *testing.T) {
+	//   ┌─────┐                 ┌─────┐
+	//   │  a  │                 │  b  │
+	//   └──┬──┘                 └──┬──┘
+	//      ├┐────────────────────►┌┤
+	//   t1 ││                     ││ t1
+	//      ├┘◄────────────────────└┤
+	//      ├┐──────────────────────►
+	//   t2 ││                      │
+	//      ├┘◄─────────────────────┤
+	g := d2graph.NewGraph(nil)
+	a := g.Root.EnsureChild([]string{"a"})
+	a.Box = geo.NewBox(nil, 100, 100)
+	a_t1 := a.EnsureChild([]string{"t1"})
+	a_t2 := a.EnsureChild([]string{"t2"})
+	b := g.Root.EnsureChild([]string{"b"})
+	b.Box = geo.NewBox(nil, 30, 30)
+	b_t1 := b.EnsureChild([]string{"t1"})
+
+	g.Edges = []*d2graph.Edge{
+		{
+			Src: a_t1,
+			Dst: b_t1,
+		}, {
+			Src: b_t1,
+			Dst: a_t1,
+		}, {
+			Src: a_t2,
+			Dst: b,
+		}, {
+			Src: b,
+			Dst: a_t2,
+		},
+	}
+
+	ctx := log.WithTB(context.Background(), t, nil)
+	Layout(ctx, g)
+
+	if a.Center().X != a_t1.Center().X {
+		t.Fatal("expected a_t1.X = a.X")
+	}
+	if a.Center().X != a_t2.Center().X {
+		t.Fatal("expected a_t2.X = a.X")
+	}
+	if b.Center().X != b_t1.Center().X {
+		t.Fatal("expected b_t1.X = b.X")
+	}
+}