fix stale *Object pointers after coreLayout

d2layouts/d2layouts.go

@@ -2,6 +2,7 @@ package d2layouts
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"sort"
 	"strings"
@@ -77,8 +78,8 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 	g.Root.Box = &geo.Box{}
 
 	// Before we can layout these nodes, we need to handle all nested diagrams first.
-	extracted := make(map[*d2graph.Object]*d2graph.Graph)
+	extracted := make(map[string]*d2graph.Graph)
-	var extractedOrder []*d2graph.Object
+	var extractedOrder []string
 
 	var constantNears []*d2graph.Graph
 	restoreOrder := SaveOrder(g)
@@ -140,6 +141,11 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 			if err != nil {
 				return err
 			}
+			// coreLayout can overwrite graph contents with newly created *Object pointers
+			// so we need to update `curr` with nestedGraph's value
+			if gi.IsConstantNear {
+				curr = nestedGraph.Root.ChildrenArray[0]
+			}
 
 			if gi.IsConstantNear {
 				curr.NearKey = nearKey
@@ -153,8 +159,10 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 				constantNears = append(constantNears, nestedGraph)
 			} else {
 				// We will restore the contents after running layout with child as the placeholder
-				extracted[curr] = nestedGraph
+				// We need to reference using ID because there may be a new object to use after coreLayout
-				extractedOrder = append(extractedOrder, curr)
+				id := curr.AbsID()
+				extracted[id] = nestedGraph
+				extractedOrder = append(extractedOrder, id)
 			}
 		} else if len(curr.ChildrenArray) > 0 {
 			queue = append(queue, curr.ChildrenArray...)
@@ -164,39 +172,52 @@ func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, co
 	// We can now run layout with accurate sizes of nested layout containers
 	// Layout according to the type of diagram
 	var err error
-	switch graphInfo.DiagramType {
+	if len(g.Objects) > 0 {
-	case GridDiagram:
+		switch graphInfo.DiagramType {
-		log.Debug(ctx, "layout grid", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
+		case GridDiagram:
-		if err = d2grid.Layout(ctx, g); err != nil {
+			log.Debug(ctx, "layout grid", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
-			return err
+			if err = d2grid.Layout(ctx, g); err != nil {
-		}
+				return err
+			}
 
-	case SequenceDiagram:
+		case SequenceDiagram:
-		log.Debug(ctx, "layout sequence", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
+			log.Debug(ctx, "layout sequence", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
-		err = d2sequence.Layout(ctx, g, coreLayout)
+			err = d2sequence.Layout(ctx, g, coreLayout)
-		if err != nil {
+			if err != nil {
-			return err
+				return err
-		}
+			}
-	default:
+		default:
-		log.Debug(ctx, "default layout", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
+			log.Debug(ctx, "default layout", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
-		err := coreLayout(ctx, g)
+			err := coreLayout(ctx, g)
-		if err != nil {
+			if err != nil {
-			return err
+				return err
+			}
 		}
 	}
 
 	if len(constantNears) > 0 {
-		err := d2near.Layout(ctx, g, constantNears)
+		err = d2near.Layout(ctx, g, constantNears)
 		if err != nil {
-			panic(err)
+			return err
 		}
 	}
 
 	// With the layout set, inject all the extracted graphs
-	for _, n := range extractedOrder {
+	for _, id := range extractedOrder {
-		nestedGraph := extracted[n]
+		nestedGraph := extracted[id]
-		InjectNested(n, nestedGraph, true)
+		// we have to find the object by ID because coreLayout can replace the Objects in graph
-		PositionNested(n, nestedGraph)
+		var obj *d2graph.Object
+		for _, o := range g.Objects {
+			if o.AbsID() == id {
+				obj = o
+				break
+			}
+		}
+		if obj == nil {
+			return fmt.Errorf("could not find object %#v after layout", id)
+		}
+		InjectNested(obj, nestedGraph, true)
+		PositionNested(obj, nestedGraph)
 	}
 
 	log.Debug(ctx, "done", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))