diff --git a/d2.go b/d2.go index d6fdfe8a2..8331ee456 100644 --- a/d2.go +++ b/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) + if layout, err := getLayout(opts); err == nil { + if err := d2sequence.Layout2(ctx, g, layout); err != nil { + return nil, err + } } else { - err = errors.New("no available layout") - } - if err != nil { 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 diff --git a/d2layouts/d2sequence/layout.go b/d2layouts/d2sequence/layout.go index f2f715a1a..9c31ac915 100644 --- a/d2layouts/d2sequence/layout.go +++ b/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:] - newObjects = append(newObjects, obj) + 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 diff --git a/d2layouts/d2sequence/layout_test.go b/d2layouts/d2sequence/layout_test.go index 1c4f24c78..2fcf61ad0 100644 --- a/d2layouts/d2sequence/layout_test.go +++ b/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,9 +324,7 @@ 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) - } + obj.TopLeft = geo.NewPoint(0, 0) } for _, edge := range g.Edges { @@ -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") } }