d2chaos.go

Documentation: oss.terrastruct.com/d2/d2chaos

     1  package d2chaos
     2  
     3  import (
     4  	"fmt"
     5  	mathrand "math/rand"
     6  	"strings"
     7  	"time"
     8  
     9  	"oss.terrastruct.com/util-go/go2"
    10  
    11  	"oss.terrastruct.com/d2/d2ast"
    12  	"oss.terrastruct.com/d2/d2format"
    13  	"oss.terrastruct.com/d2/d2graph"
    14  	"oss.terrastruct.com/d2/d2oracle"
    15  	"oss.terrastruct.com/d2/d2target"
    16  )
    17  
    18  const complexIDs = false
    19  
    20  func GenDSL(maxi int) (_ string, err error) {
    21  	gs := &dslGenState{
    22  		rand:          mathrand.New(mathrand.NewSource(time.Now().UnixNano())),
    23  		g:             d2graph.NewGraph(),
    24  		nodeShapes:    make(map[string]string),
    25  		nodeContainer: make(map[string]string),
    26  	}
    27  	gs.g.AST = &d2ast.Map{}
    28  	err = gs.gen(maxi)
    29  	if err != nil {
    30  		return "", err
    31  	}
    32  	return d2format.Format(gs.g.AST), nil
    33  }
    34  
    35  type dslGenState struct {
    36  	rand *mathrand.Rand
    37  	g    *d2graph.Graph
    38  
    39  	nodesArr      []string
    40  	nodeShapes    map[string]string
    41  	nodeContainer map[string]string
    42  }
    43  
    44  func (gs *dslGenState) gen(maxi int) error {
    45  	maxi = gs.rand.Intn(maxi) + 1
    46  
    47  	for i := 0; i < maxi; i++ {
    48  		switch gs.roll(25, 75) {
    49  		case 0:
    50  			// 25% chance of creating a new node.
    51  			err := gs.node()
    52  			if err != nil {
    53  				return err
    54  			}
    55  		case 1:
    56  			// 75% chance of connecting two random nodes with a random label.
    57  			err := gs.edge()
    58  			if err != nil {
    59  				return err
    60  			}
    61  		}
    62  	}
    63  	return nil
    64  }
    65  
    66  func (gs *dslGenState) genNode(containerID string) (string, error) {
    67  	maxLen := 8
    68  	if complexIDs {
    69  		maxLen = 32
    70  	}
    71  	nodeID := gs.randStr(maxLen, true)
    72  	if containerID != "" {
    73  		nodeID = containerID + "." + nodeID
    74  	}
    75  	var err error
    76  	gs.g, nodeID, err = d2oracle.Create(gs.g, nil, nodeID)
    77  	if err != nil {
    78  		return "", err
    79  	}
    80  	gs.nodesArr = append(gs.nodesArr, nodeID)
    81  	gs.nodeShapes[nodeID] = "square"
    82  	gs.nodeContainer[nodeID] = containerID
    83  	return nodeID, nil
    84  }
    85  
    86  func (gs *dslGenState) node() error {
    87  	containerID := ""
    88  	var err error
    89  	if gs.roll(25, 75) == 1 {
    90  		// 75% chance of creating this as a child under a container.
    91  		containerID, err = gs.randContainer()
    92  		if err != nil {
    93  			return err
    94  		}
    95  	}
    96  
    97  	nodeID, err := gs.genNode(containerID)
    98  	if err != nil {
    99  		return err
   100  	}
   101  
   102  	if gs.roll(25, 75) == 0 {
   103  		// 25% chance of adding a label.
   104  		maxLen := 8
   105  		if complexIDs {
   106  			maxLen = 256
   107  		}
   108  		gs.g, err = d2oracle.Set(gs.g, nil, nodeID, nil, go2.Pointer(gs.randStr(maxLen, false)))
   109  		if err != nil {
   110  			return err
   111  		}
   112  	}
   113  
   114  	if gs.roll(25, 75) == 1 {
   115  		// 75% chance of adding a shape.
   116  		randShape := gs.randShape()
   117  		gs.g, err = d2oracle.Set(gs.g, nil, nodeID+".shape", nil, go2.Pointer(randShape))
   118  		if err != nil {
   119  			return err
   120  		}
   121  		gs.nodeShapes[nodeID] = randShape
   122  	}
   123  
   124  	return nil
   125  }
   126  
   127  func (gs *dslGenState) edge() error {
   128  	var src string
   129  	var dst string
   130  	var err error
   131  	for {
   132  		src, err = gs.randNode()
   133  		if err != nil {
   134  			return err
   135  		}
   136  		dst, err = gs.randNode()
   137  		if err != nil {
   138  			return err
   139  		}
   140  		if gs.findOuterSequenceDiagram(src) == gs.findOuterSequenceDiagram(dst) {
   141  			break
   142  		}
   143  		err = gs.node()
   144  		if err != nil {
   145  			return err
   146  		}
   147  	}
   148  
   149  	srcArrow := "-"
   150  	if gs.randBool() {
   151  		srcArrow = "<"
   152  	}
   153  	dstArrow := "-"
   154  	if gs.randBool() {
   155  		dstArrow = ">"
   156  		if srcArrow == "<" {
   157  			dstArrow = "->"
   158  		}
   159  	}
   160  
   161  	key := fmt.Sprintf("%s %s%s %s", src, srcArrow, dstArrow, dst)
   162  	gs.g, key, err = d2oracle.Create(gs.g, nil, key)
   163  	if err != nil {
   164  		return err
   165  	}
   166  	if gs.randBool() {
   167  		maxLen := 8
   168  		if complexIDs {
   169  			maxLen = 128
   170  		}
   171  		gs.g, err = d2oracle.Set(gs.g, nil, key, nil, go2.Pointer(gs.randStr(maxLen, false)))
   172  		if err != nil {
   173  			return err
   174  		}
   175  	}
   176  	return nil
   177  }
   178  
   179  func (gs *dslGenState) randContainer() (string, error) {
   180  	containers := go2.Filter(gs.nodesArr, func(x string) bool {
   181  		shape := gs.nodeShapes[x]
   182  		return shape != "image" &&
   183  			shape != "code" &&
   184  			shape != "sql_table" &&
   185  			shape != "text" &&
   186  			shape != "class"
   187  	})
   188  	if len(containers) == 0 {
   189  		return "", nil
   190  	}
   191  	return containers[gs.rand.Intn(len(containers))], nil
   192  }
   193  
   194  func (gs *dslGenState) randNode() (string, error) {
   195  	if len(gs.nodesArr) == 0 {
   196  		return gs.genNode("")
   197  	}
   198  	return gs.nodesArr[gs.rand.Intn(len(gs.nodesArr))], nil
   199  }
   200  
   201  func (gs *dslGenState) randBool() bool {
   202  	return gs.rand.Intn(2) == 0
   203  }
   204  
   205  // TODO go back to using xrand.String, currently some incompatibility with
   206  // stuffing these strings into a script for dagre
   207  func randRune() rune {
   208  	if complexIDs {
   209  		if mathrand.Int31n(100) == 0 {
   210  			// Generate newline 1% of the time.
   211  			return '\n'
   212  		}
   213  		return mathrand.Int31n(128) + 1
   214  	} else {
   215  		return mathrand.Int31n(26) + 97
   216  	}
   217  }
   218  
   219  func (gs *dslGenState) findOuterSequenceDiagram(nodeID string) string {
   220  	for {
   221  		containerID := gs.nodeContainer[nodeID]
   222  		if containerID == "" || gs.nodeShapes[containerID] == d2target.ShapeSequenceDiagram {
   223  			return containerID
   224  		}
   225  		nodeID = containerID
   226  	}
   227  }
   228  
   229  func String(n int, exclude []rune) string {
   230  	var b strings.Builder
   231  	for i := 0; i < n; i++ {
   232  		r := randRune()
   233  		excluded := false
   234  		for _, xr := range exclude {
   235  			if r == xr {
   236  				excluded = true
   237  				break
   238  			}
   239  		}
   240  		if excluded {
   241  			i--
   242  			continue
   243  		}
   244  		b.WriteRune(r)
   245  	}
   246  	return b.String()
   247  }
   248  
   249  func (gs *dslGenState) randStr(n int, inKey bool) string {
   250  	// Underscores have semantic meaning (parent)
   251  	// Backticks are for opening and closing these strings
   252  	// Curly braces can trigger templating
   253  	// \\ triggers octal sequences
   254  	s := String(gs.rand.Intn(n), []rune{
   255  		rune('_'),
   256  		rune('`'),
   257  		rune('}'),
   258  		rune('{'),
   259  		rune('\\'),
   260  	})
   261  	as := d2ast.RawString(s, inKey)
   262  	return d2format.Format(as)
   263  }
   264  
   265  func (gs *dslGenState) randShape() string {
   266  	for {
   267  		s := shapes[gs.rand.Intn(len(shapes))]
   268  		if s != d2target.ShapeImage {
   269  			return s
   270  		}
   271  	}
   272  }
   273  
   274  func (gs *dslGenState) roll(probs ...int) int {
   275  	max := 0
   276  	for _, p := range probs {
   277  		max += p
   278  	}
   279  
   280  	n := gs.rand.Intn(max)
   281  	var acc int
   282  	for i, p := range probs {
   283  		if n >= acc && n < acc+p {
   284  			return i
   285  		}
   286  		acc += p
   287  	}
   288  
   289  	panic("d2chaos: unreachable")
   290  }
   291  
   292  var shapes = d2target.Shapes
   293
View as plain text