shift descendants

d2graph/d2graph.go

@@ -1917,3 +1917,49 @@ func (obj *Object) IterDescendants(apply func(parent, child *Object)) {
 		c.IterDescendants(apply)
 	}
 }
+
+func (obj *Object) ShiftDescendants(dx, dy float64) {
+	// also need to shift edges of descendants that are shifted
+	movedEdges := make(map[*Edge]struct{})
+	for _, e := range obj.Graph.Edges {
+		isSrcDesc := e.Src.IsDescendantOf(obj)
+		isDstDesc := e.Dst.IsDescendantOf(obj)
+
+		if isSrcDesc && isDstDesc {
+			movedEdges[e] = struct{}{}
+			for _, p := range e.Route {
+				p.X += dx
+				p.Y += dy
+			}
+		}
+	}
+
+	obj.IterDescendants(func(_, curr *Object) {
+		curr.TopLeft.X += dx
+		curr.TopLeft.Y += dy
+		for _, e := range obj.Graph.Edges {
+			if _, ok := movedEdges[e]; ok {
+				continue
+			}
+			isSrc := e.Src == curr
+			isDst := e.Dst == curr
+
+			if isSrc && isDst {
+				for _, p := range e.Route {
+					p.X += dx
+					p.Y += dy
+				}
+			} else if isSrc {
+				e.Route[0].X += dx
+				e.Route[0].Y += dy
+			} else if isDst {
+				e.Route[len(e.Route)-1].X += dx
+				e.Route[len(e.Route)-1].Y += dy
+			}
+
+			if isSrc || isDst {
+				movedEdges[e] = struct{}{}
+			}
+		}
+	})
+}

d2layouts/d2dagrelayout/layout.go

@@ -265,29 +265,6 @@ func Layout(ctx context.Context, g *d2graph.Graph, opts *ConfigurableOpts) (err
 		}
 	}
 
-	// separate loop so all descendants are placed
-	for _, obj := range g.Objects {
-		if obj.Style.ThreeDee != nil && obj.Style.ThreeDee.Value == "true" {
-			// obj.TopLeft.Y += d2target.THREE_DEE_OFFSET
-			obj.IterDescendants(func(_, child *d2graph.Object) {
-				child.TopLeft.Y += d2target.THREE_DEE_OFFSET
-			})
-			if !obj.IsContainer() {
-				obj.Height -= d2target.THREE_DEE_OFFSET
-				obj.Width -= d2target.THREE_DEE_OFFSET
-			}
-		} else if obj.Style.Multiple != nil && obj.Style.Multiple.Value == "true" {
-			// obj.TopLeft.Y += d2target.MULTIPLE_OFFSET
-			obj.IterDescendants(func(_, child *d2graph.Object) {
-				child.TopLeft.Y += d2target.MULTIPLE_OFFSET
-			})
-			if !obj.IsContainer() {
-				obj.Height -= d2target.MULTIPLE_OFFSET
-				obj.Width -= d2target.MULTIPLE_OFFSET
-			}
-		}
-	}
-
 	for i, edge := range g.Edges {
 		val, err := vm.RunString(fmt.Sprintf("JSON.stringify(g.edge(g.edges()[%d]))", i))
 		if err != nil {
@@ -549,6 +526,24 @@ func Layout(ctx context.Context, g *d2graph.Graph, opts *ConfigurableOpts) (err
 		}
 	}
 
+	// separate loop so all descendants have initial placements
+	for _, obj := range g.Objects {
+		var offset float64
+		if obj.Style.ThreeDee != nil && obj.Style.ThreeDee.Value == "true" {
+			offset = d2target.THREE_DEE_OFFSET
+		} else if obj.Style.Multiple != nil && obj.Style.Multiple.Value == "true" {
+			offset = d2target.MULTIPLE_OFFSET
+		}
+		if offset != 0 {
+			obj.TopLeft.Y += offset
+			obj.ShiftDescendants(0, offset)
+			if !obj.IsContainer() {
+				obj.Height -= offset
+				obj.Width -= offset
+			}
+		}
+	}
+
 	return nil
 }