From 4d199a4b2c697e0cfa61fb692a0085e23e345705 Mon Sep 17 00:00:00 2001
From: Nao Yonashiro <owan.orisano@gmail.com>
Date: Mon, 20 Mar 2023 04:53:43 +0900
Subject: [PATCH] feat: added the option to output opcodes as dot (#440)

---
 internal/encoder/opcode.go      | 83 +++++++++++++++++++++++++++++++++
 internal/encoder/option.go      |  1 +
 internal/encoder/vm/debug_vm.go |  6 +++
 option.go                       |  7 +++
 4 files changed, 97 insertions(+)

diff --git a/internal/encoder/opcode.go b/internal/encoder/opcode.go
index 05fc3ce..df22f55 100644
--- a/internal/encoder/opcode.go
+++ b/internal/encoder/opcode.go
@@ -1,7 +1,9 @@
 package encoder
 
 import (
+	"bytes"
 	"fmt"
+	"sort"
 	"strings"
 	"unsafe"
 
@@ -555,6 +557,87 @@ func (c *Opcode) Dump() string {
 	return strings.Join(codes, "\n")
 }
 
+func (c *Opcode) DumpDOT() string {
+	type edge struct {
+		from, to *Opcode
+		label    string
+		weight   int
+	}
+	var edges []edge
+
+	b := &bytes.Buffer{}
+	fmt.Fprintf(b, "digraph \"%p\" {\n", c.Type)
+	fmt.Fprintln(b, "mclimit=1.5;\nrankdir=TD;\nordering=out;\nnode[shape=box];")
+	for code := c; !code.IsEnd(); {
+		label := code.Op.String()
+		fmt.Fprintf(b, "\"%p\" [label=%q];\n", code, label)
+		if p := code.Next; p != nil {
+			edges = append(edges, edge{
+				from:   code,
+				to:     p,
+				label:  "Next",
+				weight: 10,
+			})
+		}
+		if p := code.NextField; p != nil {
+			edges = append(edges, edge{
+				from:   code,
+				to:     p,
+				label:  "NextField",
+				weight: 2,
+			})
+		}
+		if p := code.End; p != nil {
+			edges = append(edges, edge{
+				from:   code,
+				to:     p,
+				label:  "End",
+				weight: 1,
+			})
+		}
+		if p := code.Jmp; p != nil {
+			edges = append(edges, edge{
+				from:   code,
+				to:     p.Code,
+				label:  "Jmp",
+				weight: 1,
+			})
+		}
+
+		switch code.Op.CodeType() {
+		case CodeSliceHead:
+			code = code.Next
+		case CodeMapHead:
+			code = code.Next
+		case CodeArrayElem, CodeSliceElem:
+			code = code.End
+		case CodeMapKey:
+			code = code.End
+		case CodeMapValue:
+			code = code.Next
+		case CodeMapEnd:
+			code = code.Next
+		case CodeStructField:
+			code = code.Next
+		case CodeStructEnd:
+			code = code.Next
+		default:
+			code = code.Next
+		}
+		if code.IsEnd() {
+			fmt.Fprintf(b, "\"%p\" [label=%q];\n", code, code.Op.String())
+		}
+	}
+	sort.Slice(edges, func(i, j int) bool {
+		return edges[i].to.DisplayIdx < edges[j].to.DisplayIdx
+	})
+	for _, e := range edges {
+		fmt.Fprintf(b, "\"%p\" -> \"%p\" [label=%q][weight=%d];\n", e.from, e.to, e.label, e.weight)
+	}
+	fmt.Fprint(b, "}")
+	return b.String()
+}
+
 func newSliceHeaderCode(ctx *compileContext, typ *runtime.Type) *Opcode {
 	idx := opcodeOffset(ctx.ptrIndex)
 	ctx.incPtrIndex()
diff --git a/internal/encoder/option.go b/internal/encoder/option.go
index 82d5ce3..12c58e4 100644
--- a/internal/encoder/option.go
+++ b/internal/encoder/option.go
@@ -23,6 +23,7 @@ type Option struct {
 	ColorScheme *ColorScheme
 	Context     context.Context
 	DebugOut    io.Writer
+	DebugDOTOut io.WriteCloser
 }
 
 type EncodeFormat struct {
diff --git a/internal/encoder/vm/debug_vm.go b/internal/encoder/vm/debug_vm.go
index fbbc0de..82b6dd4 100644
--- a/internal/encoder/vm/debug_vm.go
+++ b/internal/encoder/vm/debug_vm.go
@@ -2,6 +2,7 @@ package vm
 
 import (
 	"fmt"
+	"io"
 
 	"github.com/goccy/go-json/internal/encoder"
 )
@@ -14,6 +15,11 @@ func DebugRun(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet)
 		} else {
 			code = codeSet.NoescapeKeyCode
 		}
+		if wc := ctx.Option.DebugDOTOut; wc != nil {
+			_, _ = io.WriteString(wc, code.DumpDOT())
+			wc.Close()
+			ctx.Option.DebugDOTOut = nil
+		}
 
 		if err := recover(); err != nil {
 			w := ctx.Option.DebugOut
diff --git a/option.go b/option.go
index af400a4..378031a 100644
--- a/option.go
+++ b/option.go
@@ -48,6 +48,13 @@ func DebugWith(w io.Writer) EncodeOptionFunc {
 	}
 }
 
+// DebugDOT sets the destination to write opcodes graph.
+func DebugDOT(w io.WriteCloser) EncodeOptionFunc {
+	return func(opt *EncodeOption) {
+		opt.DebugDOTOut = w
+	}
+}
+
 // Colorize add an identifier for coloring to the string of the encoded result.
 func Colorize(scheme *ColorScheme) EncodeOptionFunc {
 	return func(opt *EncodeOption) {