Merge pull request #1742 from gavin-ts/edge-router-plugin

allow plugins to route cross-graph edges

d2cli/main.go

@@ -368,6 +368,46 @@ func LayoutResolver(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plu
 	}
 }
 
+func RouterResolver(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin) func(engine string) (d2graph.RouteEdges, error) {
+	cached := make(map[string]d2graph.RouteEdges)
+	return func(engine string) (d2graph.RouteEdges, error) {
+		if c, ok := cached[engine]; ok {
+			return c, nil
+		}
+
+		plugin, err := d2plugin.FindPlugin(ctx, plugins, engine)
+		if err != nil {
+			if errors.Is(err, exec.ErrNotFound) {
+				return nil, layoutNotFound(ctx, plugins, engine)
+			}
+			return nil, err
+		}
+
+		pluginInfo, err := plugin.Info(ctx)
+		if err != nil {
+			return nil, err
+		}
+		hasRouter := false
+		for _, feat := range pluginInfo.Features {
+			if feat == d2plugin.ROUTES_EDGES {
+				hasRouter = true
+				break
+			}
+		}
+		if !hasRouter {
+			return nil, nil
+		}
+		routingPlugin, ok := plugin.(d2plugin.RoutingPlugin)
+		if !ok {
+			return nil, fmt.Errorf("plugin has routing feature but does not implement RoutingPlugin")
+		}
+
+		routeEdges := d2graph.RouteEdges(routingPlugin.RouteEdges)
+		cached[engine] = routeEdges
+		return routeEdges, nil
+	}
+}
+
 func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, fs fs.FS, layout *string, renderOpts d2svg.RenderOpts, fontFamily *d2fonts.FontFamily, animateInterval int64, inputPath, outputPath, boardPath string, bundle, forceAppendix bool, page playwright.Page) (_ []byte, written bool, _ error) {
 	start := time.Now()
 	input, err := ms.ReadPath(inputPath)
@@ -386,6 +426,7 @@ func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, fs
 		InputPath:      inputPath,
 		LayoutResolver: LayoutResolver(ctx, ms, plugins),
 		Layout:         layout,
+		RouterResolver: RouterResolver(ctx, ms, plugins),
 		FS:             fs,
 	}
 

d2exporter/export_test.go

@@ -235,7 +235,7 @@ func run(t *testing.T, tc testCase) {
 	assert.JSON(t, nil, err)
 
 	graphInfo := d2layouts.NestedGraphInfo(g.Root)
-	err = d2layouts.LayoutNested(ctx, g, graphInfo, d2dagrelayout.DefaultLayout)
+	err = d2layouts.LayoutNested(ctx, g, graphInfo, d2dagrelayout.DefaultLayout, d2layouts.DefaultRouter)
 	if err != nil {
 		t.Fatal(err)
 	}

d2graph/d2graph.go

@@ -80,6 +80,7 @@ func (g *Graph) RootBoard() *Graph {
 }
 
 type LayoutGraph func(context.Context, *Graph) error
+type RouteEdges func(context.Context, *Graph, []*Edge) error
 
 // TODO consider having different Scalar types
 // Right now we'll hold any types in Value and just convert, e.g. floats

d2layouts/d2layouts.go

@@ -76,7 +76,7 @@ func SaveOrder(g *d2graph.Graph) (restoreOrder func()) {
 	}
 }
 
-func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, coreLayout d2graph.LayoutGraph) error {
+func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, coreLayout d2graph.LayoutGraph, edgeRouter d2graph.RouteEdges) error {
 	g.Root.Box = &geo.Box{}
 
 	// Before we can layout these nodes, we need to handle all nested diagrams first.
@@ -118,7 +118,7 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 
 			// Then we layout curr as a nested graph and re-inject it
 			id := curr.AbsID()
-			err := LayoutNested(ctx, nestedGraph, GraphInfo{}, coreLayout)
+			err := LayoutNested(ctx, nestedGraph, GraphInfo{}, coreLayout, edgeRouter)
 			if err != nil {
 				return err
 			}
@@ -209,7 +209,7 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 				curr.NearKey = nil
 			}
 
-			err := LayoutNested(ctx, nestedGraph, nestedInfo, coreLayout)
+			err := LayoutNested(ctx, nestedGraph, nestedInfo, coreLayout, edgeRouter)
 			if err != nil {
 				return err
 			}
@@ -291,36 +291,61 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 		PositionNested(obj, nestedGraph)
 	}
 
-	// update map with injected objects
-	for _, o := range g.Objects {
-		idToObj[o.AbsID()] = o
-	}
+	if len(extractedEdges) > 0 {
+		// update map with injected objects
+		for _, o := range g.Objects {
+			idToObj[o.AbsID()] = o
+		}
 
-	// Restore cross-graph edges and route them
-	g.Edges = append(g.Edges, extractedEdges...)
-	for _, e := range extractedEdges {
-		// update object references
-		src, exists := idToObj[e.Src.AbsID()]
-		if !exists {
-			return fmt.Errorf("could not find object %#v after layout", e.Src.AbsID())
+		// Restore cross-graph edges and route them
+		g.Edges = append(g.Edges, extractedEdges...)
+		for _, e := range extractedEdges {
+			// update object references
+			src, exists := idToObj[e.Src.AbsID()]
+			if !exists {
+				return fmt.Errorf("could not find object %#v after layout", e.Src.AbsID())
+			}
+			e.Src = src
+			dst, exists := idToObj[e.Dst.AbsID()]
+			if !exists {
+				return fmt.Errorf("could not find object %#v after layout", e.Dst.AbsID())
+			}
+			e.Dst = dst
 		}
-		e.Src = src
-		dst, exists := idToObj[e.Dst.AbsID()]
-		if !exists {
-			return fmt.Errorf("could not find object %#v after layout", e.Dst.AbsID())
+
+		err = edgeRouter(ctx, g, extractedEdges)
+		if err != nil {
+			return err
 		}
-		e.Dst = dst
+		// need to update pointers if plugin performs edge routing
+		for _, e := range extractedEdges {
+			src, exists := idToObj[e.Src.AbsID()]
+			if !exists {
+				return fmt.Errorf("could not find object %#v after routing", e.Src.AbsID())
+			}
+			e.Src = src
+			dst, exists := idToObj[e.Dst.AbsID()]
+			if !exists {
+				return fmt.Errorf("could not find object %#v after routing", e.Dst.AbsID())
+			}
+			e.Dst = dst
+		}
+	}
 
-		// simple straight line edge routing when going across graphs
+	log.Debug(ctx, "done", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
+	return err
+}
+
+func DefaultRouter(ctx context.Context, graph *d2graph.Graph, edges []*d2graph.Edge) error {
+	for _, e := range edges {
+		// TODO replace simple straight line edge routing
 		e.Route = []*geo.Point{e.Src.Center(), e.Dst.Center()}
 		e.TraceToShape(e.Route, 0, 1)
 		if e.Label.Value != "" {
 			e.LabelPosition = go2.Pointer(label.InsideMiddleCenter.String())
 		}
 	}
-
-	log.Debug(ctx, "done", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
-	return err
+	return nil
 }
 
 func NestedGraphInfo(obj *d2graph.Object) (gi GraphInfo) {

d2lib/d2.go

@@ -25,6 +25,7 @@ type CompileOptions struct {
 	FS             fs.FS
 	MeasuredTexts  []*d2target.MText
 	Ruler          *textmeasure.Ruler
+	RouterResolver func(engine string) (d2graph.RouteEdges, error)
 	LayoutResolver func(engine string) (d2graph.LayoutGraph, error)
 
 	Layout *string
@@ -81,9 +82,13 @@ func compile(ctx context.Context, g *d2graph.Graph, compileOpts *CompileOptions,
 		if err != nil {
 			return nil, err
 		}
+		edgeRouter, err := getEdgeRouter(compileOpts)
+		if err != nil {
+			return nil, err
+		}
 
 		graphInfo := d2layouts.NestedGraphInfo(g.Root)
-		err = d2layouts.LayoutNested(ctx, g, graphInfo, coreLayout)
+		err = d2layouts.LayoutNested(ctx, g, graphInfo, coreLayout, edgeRouter)
 		if err != nil {
 			return nil, err
 		}
@@ -131,6 +136,19 @@ func getLayout(opts *CompileOptions) (d2graph.LayoutGraph, error) {
 	}
 }
 
+func getEdgeRouter(opts *CompileOptions) (d2graph.RouteEdges, error) {
+	if opts.Layout != nil && opts.RouterResolver != nil {
+		router, err := opts.RouterResolver(*opts.Layout)
+		if err != nil {
+			return nil, err
+		}
+		if router != nil {
+			return router, nil
+		}
+	}
+	return d2layouts.DefaultRouter, nil
+}
+
 // applyConfigs applies the configs read from D2 and applies it to passed in opts
 // It will only write to opt fields that are nil, as passed-in opts have precedence
 func applyConfigs(config *d2target.Config, compileOpts *CompileOptions, renderOpts *d2svg.RenderOpts) {

d2plugin/plugin.go

@@ -80,6 +80,11 @@ type Plugin interface {
 	PostProcess(context.Context, []byte) ([]byte, error)
 }
 
+type RoutingPlugin interface {
+	// RouteEdges runs the plugin's edge routing algorithm for the given edges in the input graph
+	RouteEdges(context.Context, *d2graph.Graph, []*d2graph.Edge) error
+}
+
 // PluginInfo is the current info information of a plugin.
 // note: The two fields Type and Path are not set by the plugin
 // itself but only in ListPlugins.

d2plugin/plugin_features.go

@@ -21,6 +21,9 @@ const TOP_LEFT PluginFeature = "top_left"
 // When this is true, containers can have connections to descendants
 const DESCENDANT_EDGES PluginFeature = "descendant_edges"
 
+// When this is true, the plugin also implements RoutingPlugin interface to route edges
+const ROUTES_EDGES PluginFeature = "routes_edges"
+
 func FeatureSupportCheck(info *PluginInfo, g *d2graph.Graph) error {
 	// Older version of plugin. Skip checking.
 	if info.Features == nil {