github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/graph/marshall/graph.go (about)

     1  package marshall
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"github.com/15mga/kiwi"
     7  	"github.com/15mga/kiwi/graph"
     8  	"github.com/15mga/kiwi/util"
     9  	"regexp"
    10  	"strings"
    11  )
    12  
    13  type Graph struct {
    14  }
    15  
    16  func (g *Graph) Marshall(grp graph.IGraph) []byte {
    17  	var buff util.ByteBuffer
    18  	buff.InitCap(512)
    19  	buff.WStringNoLen("\n``` mermaid")
    20  	buff.WStringNoLen("\ngraph TD")
    21  	g.marshall(grp, "", &buff)
    22  	buff.WStringNoLen("\n```")
    23  	return buff.All()
    24  }
    25  
    26  func (g *Graph) marshall(grp graph.IGraph, prefix string, buff *util.ByteBuffer) {
    27  	grp.IterNode(func(nd graph.INode) {
    28  		buff.WStringNoLen(fmt.Sprintf("\n%s  %s[%s]",
    29  			prefix, nd.Name(), nd.Comment()))
    30  	})
    31  
    32  	sp := prefix + "  "
    33  	grp.IterSubGraph(func(sg graph.ISubGraph) {
    34  		buff.WStringNoLen("\n")
    35  		buff.WStringNoLen(fmt.Sprintf("\n%ssubgraph %s", sp, sg.Name()))
    36  		g.marshall(sg, sp, buff)
    37  		buff.WStringNoLen(fmt.Sprintf("\n%send", sp))
    38  		buff.WStringNoLen("\n")
    39  	})
    40  
    41  	grp.IterLink(func(lnk graph.ILink) {
    42  		in := lnk.In().Node().Name()
    43  		ip := lnk.In().Name()
    44  		typ := lnk.In().Type()
    45  		on := lnk.Out().Node().Name()
    46  		op := lnk.Out().Name()
    47  		buff.WStringNoLen(fmt.Sprintf("\n%s  %s --> |%s->%s->%s| %s",
    48  			prefix, on, op, typ, ip, in))
    49  	})
    50  }
    51  
    52  func (g *Graph) Unmarshall(bytes []byte, grp graph.IGraph) *util.Err {
    53  	reader := bufio.NewReader(strings.NewReader(util.BytesToStr(bytes)))
    54  	ns := regexp.MustCompile(`(\[)|(])|(\()|(\))`)
    55  	ls := regexp.MustCompile(`(-->)|(\|)`)
    56  	ls2 := regexp.MustCompile(`->`)
    57  	g.unmarshal(grp, reader, ns, ls, ls2)
    58  	return nil
    59  }
    60  
    61  func (g *Graph) unmarshal(grp graph.IGraph, reader *bufio.Reader, ns, ls, ls2 *regexp.Regexp) {
    62  	linkLines := make([]string, 0, 32)
    63  	for {
    64  		lb, _, e := reader.ReadLine()
    65  		if e != nil {
    66  			break
    67  		}
    68  		line := strings.TrimSpace(string(lb))
    69  		if line == "" ||
    70  			strings.HasPrefix(line, "graph TD") ||
    71  			strings.HasPrefix(line, "subgraph") ||
    72  			strings.HasPrefix(line, "end") ||
    73  			strings.HasPrefix(line, "%%") ||
    74  			strings.HasPrefix(line, "```") {
    75  			continue
    76  		}
    77  		//连接
    78  		if strings.IndexAny(line, "-->") > -1 {
    79  			linkLines = append(linkLines, line)
    80  			continue
    81  		}
    82  		//节点
    83  		slc := ns.Split(line, -1)
    84  		slc = filterEmptyItem(slc)
    85  		nd, err := grp.AddNode(slc[0])
    86  		if err != nil {
    87  			kiwi.Error(err)
    88  			continue
    89  		}
    90  		if len(slc) == 2 {
    91  			nd.SetComment(slc[1])
    92  		}
    93  	}
    94  	for _, line := range linkLines {
    95  		slc := ls.Split(line, -1)
    96  		slc = filterEmptyItem(slc)
    97  		on := slc[0]
    98  		in := slc[2]
    99  		slc2 := ls2.Split(slc[1], -1)
   100  		slc2 = filterEmptyItem(slc2)
   101  		var ip string
   102  		op := slc2[0]
   103  		t := graph.TPoint("nil")
   104  		switch len(slc2) {
   105  		case 2:
   106  			ip = slc2[1]
   107  		case 3:
   108  			ip = slc2[2]
   109  			t = graph.TPoint(slc2[1])
   110  		default:
   111  			kiwi.Error2(util.EcParseErr, util.M{
   112  				"link": slc[1],
   113  			})
   114  		}
   115  		outNode, err := grp.GetNode(on)
   116  		if err != nil {
   117  			kiwi.Error(err)
   118  			continue
   119  		}
   120  		inNode, err := grp.GetNode(in)
   121  		if err != nil {
   122  			kiwi.Error(err)
   123  			continue
   124  		}
   125  		if !outNode.HasOut(op) {
   126  			_ = outNode.AddOut(t, op)
   127  		}
   128  		if !inNode.HasIn(ip) {
   129  			_ = inNode.AddIn(t, ip)
   130  		}
   131  		_, err = grp.Link(on, op, in, ip)
   132  		if err != nil {
   133  			kiwi.Error(err)
   134  		}
   135  	}
   136  }
   137  
   138  func filterEmptyItem(slc []string) []string {
   139  	ns := make([]string, 0, len(slc))
   140  	for _, item := range slc {
   141  		s := strings.TrimSpace(item)
   142  		if s == "" {
   143  			continue
   144  		}
   145  		ns = append(ns, s)
   146  	}
   147  	return ns
   148  }