Source file src/oss.terrastruct.com/d2/d2layouts/d2layouts.go

Documentation: oss.terrastruct.com/d2/d2layouts

     1  package d2layouts
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math"
     7  	"sort"
     8  	"strings"
     9  
    10  	"cdr.dev/slog"
    11  
    12  	"oss.terrastruct.com/d2/d2graph"
    13  	"oss.terrastruct.com/d2/d2layouts/d2grid"
    14  	"oss.terrastruct.com/d2/d2layouts/d2near"
    15  	"oss.terrastruct.com/d2/d2layouts/d2sequence"
    16  	"oss.terrastruct.com/d2/lib/geo"
    17  	"oss.terrastruct.com/d2/lib/label"
    18  	"oss.terrastruct.com/d2/lib/log"
    19  	"oss.terrastruct.com/util-go/go2"
    20  )
    21  
    22  type DiagramType string
    23  
    24  // a grid diagram at a constant near is
    25  const (
    26  	DefaultGraphType  DiagramType = ""
    27  	ConstantNearGraph DiagramType = "constant-near"
    28  	GridDiagram       DiagramType = "grid-diagram"
    29  	SequenceDiagram   DiagramType = "sequence-diagram"
    30  )
    31  
    32  type GraphInfo struct {
    33  	IsConstantNear bool
    34  	DiagramType    DiagramType
    35  }
    36  
    37  func (gi GraphInfo) isDefault() bool {
    38  	return !gi.IsConstantNear && gi.DiagramType == DefaultGraphType
    39  }
    40  
    41  func SaveChildrenOrder(container *d2graph.Object) (restoreOrder func()) {
    42  	objectOrder := make(map[string]int, len(container.ChildrenArray))
    43  	for i, obj := range container.ChildrenArray {
    44  		objectOrder[obj.AbsID()] = i
    45  	}
    46  	return func() {
    47  		sort.SliceStable(container.ChildrenArray, func(i, j int) bool {
    48  			return objectOrder[container.ChildrenArray[i].AbsID()] < objectOrder[container.ChildrenArray[j].AbsID()]
    49  		})
    50  	}
    51  }
    52  
    53  func SaveOrder(g *d2graph.Graph) (restoreOrder func()) {
    54  	objectOrder := make(map[string]int, len(g.Objects))
    55  	for i, obj := range g.Objects {
    56  		objectOrder[obj.AbsID()] = i
    57  	}
    58  	edgeOrder := make(map[string]int, len(g.Edges))
    59  	for i, edge := range g.Edges {
    60  		edgeOrder[edge.AbsID()] = i
    61  	}
    62  	restoreRootOrder := SaveChildrenOrder(g.Root)
    63  	return func() {
    64  		sort.SliceStable(g.Objects, func(i, j int) bool {
    65  			return objectOrder[g.Objects[i].AbsID()] < objectOrder[g.Objects[j].AbsID()]
    66  		})
    67  		sort.SliceStable(g.Edges, func(i, j int) bool {
    68  			iIndex, iHas := edgeOrder[g.Edges[i].AbsID()]
    69  			jIndex, jHas := edgeOrder[g.Edges[j].AbsID()]
    70  			if iHas && jHas {
    71  				return iIndex < jIndex
    72  			}
    73  			return iHas
    74  		})
    75  		restoreRootOrder()
    76  	}
    77  }
    78  
    79  func LayoutNested(ctx context.Context, g *d2graph.Graph, graphInfo GraphInfo, coreLayout d2graph.LayoutGraph, edgeRouter d2graph.RouteEdges) error {
    80  	g.Root.Box = &geo.Box{}
    81  
    82  	// Before we can layout these nodes, we need to handle all nested diagrams first.
    83  	extracted := make(map[string]*d2graph.Graph)
    84  	var extractedOrder []string
    85  	var extractedEdges []*d2graph.Edge
    86  	var extractedEdgeIDs []edgeIDs
    87  
    88  	var constantNears []*d2graph.Graph
    89  	restoreOrder := SaveOrder(g)
    90  	defer restoreOrder()
    91  
    92  	// Iterate top-down from Root so all nested diagrams can process their own contents
    93  	queue := make([]*d2graph.Object, 0, len(g.Root.ChildrenArray))
    94  	queue = append(queue, g.Root.ChildrenArray...)
    95  
    96  	for len(queue) > 0 {
    97  		curr := queue[0]
    98  		queue = queue[1:]
    99  
   100  		isGridCellContainer := graphInfo.DiagramType == GridDiagram &&
   101  			curr.IsContainer() && curr.Parent == g.Root
   102  		gi := NestedGraphInfo(curr)
   103  
   104  		if isGridCellContainer && gi.isDefault() {
   105  			// if we are in a grid diagram, and our children have descendants
   106  			// we need to run layout on them first, even if they are not special diagram types
   107  
   108  			// First we extract the grid cell container as a nested graph with includeSelf=true
   109  			// resulting in externalEdges=[A, C] and nestedGraph.Edges=[B]
   110  			// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
   111  			// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐                        │
   112  			// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │                        │
   113  			// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    │                        │
   114  			// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │                        │
   115  			// │ │    │      │  └──┘    └──┘ │ │    │ │    │                        │
   116  			// │ └────┘      └───────────────┘ │    │ └────┘                        │
   117  			// └───────────────────────────────┘    └───────────────────────────────┘
   118  			nestedGraph, externalEdges, externalEdgeIDs := ExtractSubgraph(curr, true)
   119  
   120  			// Then we layout curr as a nested graph and re-inject it
   121  			id := curr.AbsID()
   122  			err := LayoutNested(ctx, nestedGraph, GraphInfo{}, coreLayout, edgeRouter)
   123  			if err != nil {
   124  				return err
   125  			}
   126  			InjectNested(g.Root, nestedGraph, false)
   127  			g.Edges = append(g.Edges, externalEdges...)
   128  			restoreOrder()
   129  
   130  			// layout can replace Objects so we need to update the references we are holding onto (curr + externalEdges)
   131  			idToObj := make(map[string]*d2graph.Object)
   132  			for _, o := range g.Objects {
   133  				idToObj[o.AbsID()] = o
   134  			}
   135  			lookup := func(idStr string) (*d2graph.Object, error) {
   136  				o, exists := idToObj[idStr]
   137  				if !exists {
   138  					return nil, fmt.Errorf("could not find object %#v after layout", idStr)
   139  				}
   140  				return o, nil
   141  			}
   142  			curr, err = lookup(id)
   143  			if err != nil {
   144  				return err
   145  			}
   146  			for i, e := range externalEdges {
   147  				src, err := lookup(externalEdgeIDs[i].srcID)
   148  				if err != nil {
   149  					return err
   150  				}
   151  				e.Src = src
   152  				dst, err := lookup(externalEdgeIDs[i].dstID)
   153  				if err != nil {
   154  					return err
   155  				}
   156  				e.Dst = dst
   157  			}
   158  
   159  			// position nested graph (excluding curr) relative to curr
   160  			dx := 0 - curr.TopLeft.X
   161  			dy := 0 - curr.TopLeft.Y
   162  			for _, o := range nestedGraph.Objects {
   163  				if o == curr {
   164  					continue
   165  				}
   166  				o.TopLeft.X += dx
   167  				o.TopLeft.Y += dy
   168  			}
   169  			for _, e := range nestedGraph.Edges {
   170  				e.Move(dx, dy)
   171  			}
   172  
   173  			// Then after re-injecting everything, we extract curr with includeSelf=false,
   174  			// and externalEdges=[C], nestedGraph.Edges=[B], and graph.Edges=[A].
   175  			// This will leave A in the graph to be routed by grid layout, and C will have cross-graph edge routing
   176  			// Note: currently grid layout's cell-cell edge routing, and cross-graph edge routing behave the same,
   177  			// but these are simple placeholder routings and they may be different in the future
   178  			// ┌grid(g.Root)───────────────────┐    ┌grid(g.Root)───────────────────┐
   179  			// │ ┌────┐      ┌curr───────────┐ │    │ ┌────┐      ┌curr───────────┐ │
   180  			// │ │    │      │  ┌──┐    ┌──┐ │ │    │ │    │      │               │ │
   181  			// │ │    ├──A──►│  │  ├─B─►│  │ │ │ => │ │    ├──A──►│               │ │
   182  			// │ │    ├──────┼C─┼──┼───►│  │ │ │    │ │    │      │               │ │
   183  			// │ │    │      │  └──┘    └──┘ │ │    │ │    │      │               │ │
   184  			// │ └────┘      └───────────────┘ │    │ └────┘      └───────────────┘ │
   185  			// └───────────────────────────────┘    └───────────────────────────────┘
   186  			nestedGraph, externalEdges, externalEdgeIDs = ExtractSubgraph(curr, false)
   187  			extractedEdges = append(extractedEdges, externalEdges...)
   188  			extractedEdgeIDs = append(extractedEdgeIDs, externalEdgeIDs...)
   189  
   190  			extracted[id] = nestedGraph
   191  			extractedOrder = append(extractedOrder, id)
   192  			continue
   193  		}
   194  
   195  		if !gi.isDefault() {
   196  			// empty grid or sequence can have 0 objects..
   197  			if !gi.IsConstantNear && len(curr.Children) == 0 {
   198  				continue
   199  			}
   200  
   201  			// There is a nested diagram here, so extract its contents and process in the same way
   202  			nestedGraph, externalEdges, externalEdgeIDs := ExtractSubgraph(curr, gi.IsConstantNear)
   203  			extractedEdges = append(extractedEdges, externalEdges...)
   204  			extractedEdgeIDs = append(extractedEdgeIDs, externalEdgeIDs...)
   205  
   206  			log.Info(ctx, "layout nested", slog.F("level", curr.Level()), slog.F("child", curr.AbsID()), slog.F("gi", gi))
   207  			nestedInfo := gi
   208  			nearKey := curr.NearKey
   209  			if gi.IsConstantNear {
   210  				// layout nested as a non-near
   211  				nestedInfo = GraphInfo{}
   212  				curr.NearKey = nil
   213  			}
   214  
   215  			err := LayoutNested(ctx, nestedGraph, nestedInfo, coreLayout, edgeRouter)
   216  			if err != nil {
   217  				return err
   218  			}
   219  			// coreLayout can overwrite graph contents with newly created *Object pointers
   220  			// so we need to update `curr` with nestedGraph's value
   221  			if gi.IsConstantNear {
   222  				curr = nestedGraph.Root.ChildrenArray[0]
   223  			}
   224  
   225  			if gi.IsConstantNear {
   226  				curr.NearKey = nearKey
   227  			} else {
   228  				FitToGraph(curr, nestedGraph, geo.Spacing{})
   229  				curr.TopLeft = geo.NewPoint(0, 0)
   230  			}
   231  
   232  			if gi.IsConstantNear {
   233  				// near layout will inject these nestedGraphs
   234  				constantNears = append(constantNears, nestedGraph)
   235  			} else {
   236  				// We will restore the contents after running layout with child as the placeholder
   237  				// We need to reference using ID because there may be a new object to use after coreLayout
   238  				id := curr.AbsID()
   239  				extracted[id] = nestedGraph
   240  				extractedOrder = append(extractedOrder, id)
   241  			}
   242  		} else if len(curr.ChildrenArray) > 0 {
   243  			queue = append(queue, curr.ChildrenArray...)
   244  		}
   245  	}
   246  
   247  	// We can now run layout with accurate sizes of nested layout containers
   248  	// Layout according to the type of diagram
   249  	var err error
   250  	if len(g.Objects) > 0 {
   251  		switch graphInfo.DiagramType {
   252  		case GridDiagram:
   253  			log.Debug(ctx, "layout grid", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
   254  			if err = d2grid.Layout(ctx, g); err != nil {
   255  				return err
   256  			}
   257  
   258  		case SequenceDiagram:
   259  			log.Debug(ctx, "layout sequence", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
   260  			err = d2sequence.Layout(ctx, g, coreLayout)
   261  			if err != nil {
   262  				return err
   263  			}
   264  		default:
   265  			log.Debug(ctx, "default layout", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
   266  			err := coreLayout(ctx, g)
   267  			if err != nil {
   268  				return err
   269  			}
   270  		}
   271  	}
   272  
   273  	if len(constantNears) > 0 {
   274  		err = d2near.Layout(ctx, g, constantNears)
   275  		if err != nil {
   276  			return err
   277  		}
   278  	}
   279  
   280  	idToObj := make(map[string]*d2graph.Object)
   281  	for _, o := range g.Objects {
   282  		idToObj[o.AbsID()] = o
   283  	}
   284  
   285  	// With the layout set, inject all the extracted graphs
   286  	for _, id := range extractedOrder {
   287  		nestedGraph := extracted[id]
   288  		// we have to find the object by ID because coreLayout can replace the Objects in graph
   289  		obj, exists := idToObj[id]
   290  		if !exists {
   291  			return fmt.Errorf("could not find object %#v after layout", id)
   292  		}
   293  		InjectNested(obj, nestedGraph, true)
   294  		PositionNested(obj, nestedGraph)
   295  	}
   296  
   297  	if len(extractedEdges) > 0 {
   298  		// update map with injected objects
   299  		for _, o := range g.Objects {
   300  			idToObj[o.AbsID()] = o
   301  		}
   302  
   303  		// Restore cross-graph edges and route them
   304  		g.Edges = append(g.Edges, extractedEdges...)
   305  		for i, e := range extractedEdges {
   306  			ids := extractedEdgeIDs[i]
   307  			// update object references
   308  			src, exists := idToObj[ids.srcID]
   309  			if !exists {
   310  				return fmt.Errorf("could not find object %#v after layout", ids.srcID)
   311  			}
   312  			e.Src = src
   313  			dst, exists := idToObj[ids.dstID]
   314  			if !exists {
   315  				return fmt.Errorf("could not find object %#v after layout", ids.dstID)
   316  			}
   317  			e.Dst = dst
   318  		}
   319  
   320  		err = edgeRouter(ctx, g, extractedEdges)
   321  		if err != nil {
   322  			return err
   323  		}
   324  		// need to update pointers if plugin performs edge routing
   325  		for i, e := range extractedEdges {
   326  			ids := extractedEdgeIDs[i]
   327  			src, exists := idToObj[ids.srcID]
   328  			if !exists {
   329  				return fmt.Errorf("could not find object %#v after routing", ids.srcID)
   330  			}
   331  			e.Src = src
   332  			dst, exists := idToObj[ids.dstID]
   333  			if !exists {
   334  				return fmt.Errorf("could not find object %#v after routing", ids.dstID)
   335  			}
   336  			e.Dst = dst
   337  		}
   338  	}
   339  
   340  	log.Debug(ctx, "done", slog.F("rootlevel", g.RootLevel), slog.F("shapes", g.PrintString()))
   341  	return err
   342  }
   343  
   344  func DefaultRouter(ctx context.Context, graph *d2graph.Graph, edges []*d2graph.Edge) error {
   345  	for _, e := range edges {
   346  		// TODO replace simple straight line edge routing
   347  		e.Route = []*geo.Point{e.Src.Center(), e.Dst.Center()}
   348  		e.TraceToShape(e.Route, 0, 1)
   349  		if e.Label.Value != "" {
   350  			e.LabelPosition = go2.Pointer(label.InsideMiddleCenter.String())
   351  		}
   352  	}
   353  	return nil
   354  }
   355  
   356  func NestedGraphInfo(obj *d2graph.Object) (gi GraphInfo) {
   357  	if obj.Graph.RootLevel == 0 && obj.IsConstantNear() {
   358  		gi.IsConstantNear = true
   359  	}
   360  	if obj.IsSequenceDiagram() {
   361  		gi.DiagramType = SequenceDiagram
   362  	} else if obj.IsGridDiagram() {
   363  		gi.DiagramType = GridDiagram
   364  	}
   365  	return gi
   366  }
   367  
   368  type edgeIDs struct {
   369  	srcID, dstID string
   370  }
   371  
   372  func ExtractSubgraph(container *d2graph.Object, includeSelf bool) (nestedGraph *d2graph.Graph, externalEdges []*d2graph.Edge, externalEdgeIDs []edgeIDs) {
   373  	// includeSelf: when we have a constant near or a grid cell that is a container,
   374  	// we want to include itself in the nested graph, not just its descendants,
   375  	nestedGraph = d2graph.NewGraph()
   376  	nestedGraph.RootLevel = int(container.Level())
   377  	if includeSelf {
   378  		nestedGraph.RootLevel--
   379  	}
   380  	nestedGraph.Root.Attributes = container.Attributes
   381  	nestedGraph.Root.Box = &geo.Box{}
   382  	nestedGraph.Root.LabelPosition = container.LabelPosition
   383  	nestedGraph.Root.IconPosition = container.IconPosition
   384  
   385  	isNestedObject := func(obj *d2graph.Object) bool {
   386  		if includeSelf {
   387  			return obj.IsDescendantOf(container)
   388  		}
   389  		return obj.Parent.IsDescendantOf(container)
   390  	}
   391  
   392  	// separate out nested edges
   393  	g := container.Graph
   394  	remainingEdges := make([]*d2graph.Edge, 0, len(g.Edges))
   395  	for _, edge := range g.Edges {
   396  		srcIsNested := isNestedObject(edge.Src)
   397  		if d2sequence.IsLifelineEnd(edge.Dst) {
   398  			// special handling for lifelines since their edge.Dst is a special Object
   399  			if srcIsNested {
   400  				nestedGraph.Edges = append(nestedGraph.Edges, edge)
   401  			} else {
   402  				remainingEdges = append(remainingEdges, edge)
   403  			}
   404  			continue
   405  		}
   406  		dstIsNested := isNestedObject(edge.Dst)
   407  		if srcIsNested && dstIsNested {
   408  			nestedGraph.Edges = append(nestedGraph.Edges, edge)
   409  		} else if srcIsNested || dstIsNested {
   410  			externalEdges = append(externalEdges, edge)
   411  			// we need these AbsIDs when reconnecting since parent references may become stale
   412  			externalEdgeIDs = append(externalEdgeIDs, edgeIDs{
   413  				srcID: edge.Src.AbsID(),
   414  				dstID: edge.Dst.AbsID(),
   415  			})
   416  		} else {
   417  			remainingEdges = append(remainingEdges, edge)
   418  		}
   419  	}
   420  	g.Edges = remainingEdges
   421  
   422  	// separate out nested objects
   423  	remainingObjects := make([]*d2graph.Object, 0, len(g.Objects))
   424  	for _, obj := range g.Objects {
   425  		if isNestedObject(obj) {
   426  			nestedGraph.Objects = append(nestedGraph.Objects, obj)
   427  		} else {
   428  			remainingObjects = append(remainingObjects, obj)
   429  		}
   430  	}
   431  	g.Objects = remainingObjects
   432  
   433  	// update object and new root references
   434  	for _, o := range nestedGraph.Objects {
   435  		o.Graph = nestedGraph
   436  	}
   437  
   438  	if includeSelf {
   439  		// remove container parent's references
   440  		if container.Parent != nil {
   441  			container.Parent.RemoveChild(container)
   442  		}
   443  
   444  		// set root references
   445  		nestedGraph.Root.ChildrenArray = []*d2graph.Object{container}
   446  		container.Parent = nestedGraph.Root
   447  		nestedGraph.Root.Children[strings.ToLower(container.ID)] = container
   448  	} else {
   449  		// set root references
   450  		nestedGraph.Root.ChildrenArray = append(nestedGraph.Root.ChildrenArray, container.ChildrenArray...)
   451  		for _, child := range container.ChildrenArray {
   452  			child.Parent = nestedGraph.Root
   453  			nestedGraph.Root.Children[strings.ToLower(child.ID)] = child
   454  		}
   455  
   456  		// remove container's references
   457  		for k := range container.Children {
   458  			delete(container.Children, k)
   459  		}
   460  		container.ChildrenArray = nil
   461  	}
   462  
   463  	return nestedGraph, externalEdges, externalEdgeIDs
   464  }
   465  
   466  func InjectNested(container *d2graph.Object, nestedGraph *d2graph.Graph, isRoot bool) {
   467  	g := container.Graph
   468  	for _, obj := range nestedGraph.Root.ChildrenArray {
   469  		obj.Parent = container
   470  		if container.Children == nil {
   471  			container.Children = make(map[string]*d2graph.Object)
   472  		}
   473  		container.Children[strings.ToLower(obj.ID)] = obj
   474  		container.ChildrenArray = append(container.ChildrenArray, obj)
   475  	}
   476  	for _, obj := range nestedGraph.Objects {
   477  		obj.Graph = g
   478  	}
   479  	g.Objects = append(g.Objects, nestedGraph.Objects...)
   480  	g.Edges = append(g.Edges, nestedGraph.Edges...)
   481  
   482  	if isRoot {
   483  		if nestedGraph.Root.LabelPosition != nil {
   484  			container.LabelPosition = nestedGraph.Root.LabelPosition
   485  		}
   486  		if nestedGraph.Root.IconPosition != nil {
   487  			container.IconPosition = nestedGraph.Root.IconPosition
   488  		}
   489  		container.Attributes = nestedGraph.Root.Attributes
   490  	}
   491  }
   492  
   493  func PositionNested(container *d2graph.Object, nestedGraph *d2graph.Graph) {
   494  	// tl, _ := boundingBox(nestedGraph)
   495  	// Note: assumes nestedGraph's layout has contents positioned relative to 0,0
   496  	dx := container.TopLeft.X //- tl.X
   497  	dy := container.TopLeft.Y //- tl.Y
   498  	if dx == 0 && dy == 0 {
   499  		return
   500  	}
   501  	for _, o := range nestedGraph.Objects {
   502  		o.TopLeft.X += dx
   503  		o.TopLeft.Y += dy
   504  	}
   505  	for _, e := range nestedGraph.Edges {
   506  		e.Move(dx, dy)
   507  	}
   508  }
   509  
   510  func boundingBox(g *d2graph.Graph) (tl, br *geo.Point) {
   511  	if len(g.Objects) == 0 {
   512  		return geo.NewPoint(0, 0), geo.NewPoint(0, 0)
   513  	}
   514  	tl = geo.NewPoint(math.Inf(1), math.Inf(1))
   515  	br = geo.NewPoint(math.Inf(-1), math.Inf(-1))
   516  
   517  	for _, obj := range g.Objects {
   518  		if obj.TopLeft == nil {
   519  			panic(obj.AbsID())
   520  		}
   521  		tl.X = math.Min(tl.X, obj.TopLeft.X)
   522  		tl.Y = math.Min(tl.Y, obj.TopLeft.Y)
   523  		br.X = math.Max(br.X, obj.TopLeft.X+obj.Width)
   524  		br.Y = math.Max(br.Y, obj.TopLeft.Y+obj.Height)
   525  	}
   526  
   527  	return tl, br
   528  }
   529  
   530  func FitToGraph(container *d2graph.Object, nestedGraph *d2graph.Graph, padding geo.Spacing) {
   531  	var width, height float64
   532  	width = nestedGraph.Root.Width
   533  	height = nestedGraph.Root.Height
   534  	if width == 0 || height == 0 {
   535  		tl, br := boundingBox(nestedGraph)
   536  		width = br.X - tl.X
   537  		height = br.Y - tl.Y
   538  	}
   539  	container.Width = padding.Left + width + padding.Right
   540  	container.Height = padding.Top + height + padding.Bottom
   541  }
   542  

View as plain text