actual fix: correct stale object references

d2layouts/d2layouts.go

@@ -103,77 +103,63 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 		if isGridCellContainer && gi.isDefault() {
 			// if we are in a grid diagram, and our children have descendants
 			// we need to run layout on them first, even if they are not special diagram types
+
+			// First we extract the grid cell container as a nested graph with includeSelf=true
+			// resulting in externalEdges=[A, C] and nestedGraph.Edges=[B]
+			// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
+			// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐                        │
+			// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │                        │
+			// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    │                        │
+			// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │                        │
+			// │ │    │      │  └──┘    └──┘ │ │    │ │    │                        │
+			// │ └────┘      └───────────────┘ │    │ └────┘                        │
+			// └───────────────────────────────┘    └───────────────────────────────┘
 			nestedGraph, externalEdges := ExtractSubgraph(curr, true)
+
+			// Then we layout curr as a nested graph and re-inject it
 			id := curr.AbsID()
 			err := LayoutNested(ctx, nestedGraph, GraphInfo{}, coreLayout)
 			if err != nil {
 				return err
 			}
-
 			InjectNested(g.Root, nestedGraph, false)
-			for _, e := range externalEdges {
+			g.Edges = append(g.Edges, externalEdges...)
-				if e.Src.AbsID() == id || e.Dst.AbsID() == id {
-					// We need to keep this edge extracted to route after grid layout.
-					// When we extracted curr(grid cell container) externalEdges=[A, C] and nestedGraph.Edges=[B]
-					// With includeSelf=true
-					// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
-					// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐                        │
-					// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │                        │
-					// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    │                        │
-					// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │                        │
-					// │ │    │      │  └──┘    └──┘ │ │    │ │    │                        │
-					// │ └────┘      └───────────────┘ │    │ └────┘                        │
-					// └───────────────────────────────┘    └───────────────────────────────┘
-					//
-					// If we add back all external edges, when we extract curr with includeSelf=false,
-					// externalEdges would be [C], nestedGraph.Edges=[B], and graph.Edges=[A].
-					// This would incorrectly leave A in the graph and it wouldn't be routed after grid layout
-					// With includeSelf=false
-					// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
-					// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐      ┌curr───────────┐ │
-					// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │      │               │ │
-					// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    ├──A──►│               │ │
-					// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │      │               │ │
-					// │ │    │      │  └──┘    └──┘ │ │    │ │    │      │               │ │
-					// │ └────┘      └───────────────┘ │    │ └────┘      └───────────────┘ │
-					// └───────────────────────────────┘    └───────────────────────────────┘
-					//
-					// Therefore we only add [B, C] since they will be correctly extracted with includeSelf=false
-					// as externalEdges=[C] and nestedGraph.Edges=[B].
-					// With includeSelf=false
-					// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
-					// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐      ┌curr───────────┐ │
-					// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │      │               │ │
-					// │ │    │      │  │  ├─B─►│  │ │ │ => │ │    │      │               │ │
-					// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │      │               │ │
-					// │ │    │      │  └──┘    └──┘ │ │    │ │    │      │               │ │
-					// │ └────┘      └───────────────┘ │    │ └────┘      └───────────────┘ │
-					// └───────────────────────────────┘    └───────────────────────────────┘
-					extractedEdges = append(extractedEdges, e)
-				} else {
-					g.Edges = append(g.Edges, e)
-				}
-			}
 			restoreOrder()
 
-			// need to update curr *Object incase layout changed it
+			// layout can replace Objects so we need to update the references we are holding onto (curr + externalEdges)
-			var obj *d2graph.Object
+			idToObj := make(map[string]*d2graph.Object)
 			for _, o := range g.Objects {
-				if o.AbsID() == id {
+				idToObj[o.AbsID()] = o
-					obj = o
+			}
-					break
+			lookup := func(idStr string) (*d2graph.Object, error) {
+				o, exists := idToObj[idStr]
+				if !exists {
+					return nil, fmt.Errorf("could not find object %#v after layout", idStr)
 				}
+				return o, nil
 			}
-			if obj == nil {
+			curr, err = lookup(id)
-				return fmt.Errorf("could not find object %#v after layout", id)
+			if err != nil {
+				return err
+			}
+			for _, e := range externalEdges {
+				src, err := lookup(e.Src.AbsID())
+				if err != nil {
+					return err
+				}
+				e.Src = src
+				dst, err := lookup(e.Dst.AbsID())
+				if err != nil {
+					return err
+				}
+				e.Dst = dst
 			}
-			curr = obj
 
 			// position nested graph (excluding curr) relative to curr
 			dx := 0 - curr.TopLeft.X
 			dy := 0 - curr.TopLeft.Y
 			for _, o := range nestedGraph.Objects {
-				if o.AbsID() == curr.AbsID() {
+				if o == curr {
 					continue
 				}
 				o.TopLeft.X += dx
@@ -183,7 +169,19 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 				e.Move(dx, dy)
 			}
 
-			// now we keep the descendants out until after grid layout
+			// Then after re-injecting everything, we extract curr with includeSelf=false,
+			// and externalEdges=[C], nestedGraph.Edges=[B], and graph.Edges=[A].
+			// This will leave A in the graph to be routed by grid layout, and C will have cross-graph edge routing
+			// Note: currently grid layout's cell-cell edge routing, and cross-graph edge routing behave the same,
+			// but these are simple placeholder routings and they may be different in the future
+			// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
+			// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐      ┌curr───────────┐ │
+			// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │      │               │ │
+			// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    ├──A──►│               │ │
+			// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │      │               │ │
+			// │ │    │      │  └──┘    └──┘ │ │    │ │    │      │               │ │
+			// │ └────┘      └───────────────┘ │    │ └────┘      └───────────────┘ │
+			// └───────────────────────────────┘    └───────────────────────────────┘
 			nestedGraph, externalEdges = ExtractSubgraph(curr, false)
 			extractedEdges = append(extractedEdges, externalEdges...)