Layout sequence diagrams before running the layout engine

d2.go

@@ -9,6 +9,7 @@ import (
 	"oss.terrastruct.com/d2/d2compiler"
 	"oss.terrastruct.com/d2/d2exporter"
 	"oss.terrastruct.com/d2/d2graph"
+	"oss.terrastruct.com/d2/d2layouts/d2sequence"
 	"oss.terrastruct.com/d2/d2renderers/textmeasure"
 	"oss.terrastruct.com/d2/d2target"
 )
@@ -39,14 +40,11 @@ func Compile(ctx context.Context, input string, opts *CompileOptions) (*d2target
 		return nil, err
 	}
 
-	if opts.Layout != nil {
-		err = opts.Layout(ctx, g)
-	} else if os.Getenv("D2_LAYOUT") == "dagre" && dagreLayout != nil {
-		err = dagreLayout(ctx, g)
-	} else {
-		err = errors.New("no available layout")
+	if layout, err := getLayout(opts); err == nil {
+		if err := d2sequence.Layout2(ctx, g, layout); err != nil {
+			return nil, err
 		}
-	if err != nil {
+	} else {
 		return nil, err
 	}
 
@@ -54,5 +52,15 @@ func Compile(ctx context.Context, input string, opts *CompileOptions) (*d2target
 	return diagram, err
 }
 
+func getLayout(opts *CompileOptions) (func(context.Context, *d2graph.Graph) error, error) {
+	if opts.Layout != nil {
+		return opts.Layout, nil
+	} else if os.Getenv("D2_LAYOUT") == "dagre" && dagreLayout != nil {
+		return dagreLayout, nil
+	} else {
+		return nil, errors.New("no available layout")
+	}
+}
+
 // See c.go
 var dagreLayout func(context.Context, *d2graph.Graph) error

d2layouts/d2sequence/layout.go

@@ -16,6 +16,9 @@ import (
 )
 
 func Layout2(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Context, g *d2graph.Graph) error) error {
+	oldObjects := g.Objects
+	oldEdges := g.Edges
+
 	// 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{})
@@ -30,7 +33,9 @@ func Layout2(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Cont
 		obj := queue[0]
 		queue = queue[1:]
 
+		if obj != g.Root {
 			newObjects = append(newObjects, obj)
+		}
 		if obj.Attributes.Shape.Value == d2target.ShapeSequenceDiagram {
 			// TODO: should update obj.References too?
 
@@ -76,22 +81,17 @@ func Layout2(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Cont
 	}
 
 	// restores objects & edges
-	for edge := range edgesToRemove {
-		g.Edges = append(g.Edges, edge)
-	}
+	g.Edges = oldEdges
+	g.Objects = oldObjects
 	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)
-		}
+		g.Edges = append(g.Edges, sequenceDiagrams[obj].lifelines...)
 	}
 
 	return nil

d2layouts/d2sequence/layout_test.go

@@ -296,8 +296,8 @@ 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 {
+		// 3 because it replaces all `container` children with a rectangle for layout
+		if len(g.Objects) != 3 {
 			t.Fatal("expected only diagram objects for layout")
 		}
 		for _, obj := range g.Objects {
@@ -324,10 +324,8 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 
 		// 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)}
@@ -346,7 +344,7 @@ func TestNestedSequenceDiagrams(t *testing.T) {
 	}
 
 	for _, obj := range g.Objects {
-		if obj != g.Root && obj.TopLeft == nil {
+		if obj.TopLeft == nil {
 			t.Fatal("expected to have placed all objects")
 		}
 	}