From adb139a4fd62f8bcebab9ba9c7d4d7d6e61e13ca Mon Sep 17 00:00:00 2001
From: Anton Medvedev <anton@medv.io>
Date: Thu, 29 Feb 2024 20:51:34 +0100
Subject: [PATCH] Add spans

---
 compiler/compiler.go | 29 +++++++++++++++++++++++++++++
 conf/config.go       |  1 +
 vm/opcodes.go        |  2 ++
 vm/program.go        |  9 +++++++++
 vm/utils.go          | 13 +++++++++++++
 vm/vm.go             |  9 +++++++++
 6 files changed, 63 insertions(+)

diff --git a/compiler/compiler.go b/compiler/compiler.go
index a4f189e6..808b53c9 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -50,6 +50,11 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 		}
 	}
 
+	var span *Span
+	if len(c.spans) > 0 {
+		span = c.spans[0]
+	}
+
 	program = NewProgram(
 		tree.Source,
 		tree.Node,
@@ -60,6 +65,7 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 		c.arguments,
 		c.functions,
 		c.debugInfo,
+		span,
 	)
 	return
 }
@@ -76,6 +82,7 @@ type compiler struct {
 	functionsIndex map[string]int
 	debugInfo      map[string]string
 	nodes          []ast.Node
+	spans          []*Span
 	chains         [][]int
 	arguments      []int
 }
@@ -193,6 +200,28 @@ func (c *compiler) compile(node ast.Node) {
 		c.nodes = c.nodes[:len(c.nodes)-1]
 	}()
 
+	if c.config != nil && c.config.Profile {
+		span := &Span{
+			Name:       reflect.TypeOf(node).String(),
+			Expression: node.String(),
+		}
+		if len(c.spans) > 0 {
+			prev := c.spans[len(c.spans)-1]
+			prev.Children = append(prev.Children, span)
+		}
+		c.spans = append(c.spans, span)
+		defer func() {
+			if len(c.spans) > 1 {
+				c.spans = c.spans[:len(c.spans)-1]
+			}
+		}()
+
+		c.emit(OpProfileStart, c.addConstant(span))
+		defer func() {
+			c.emit(OpProfileEnd, c.addConstant(span))
+		}()
+	}
+
 	switch n := node.(type) {
 	case *ast.NilNode:
 		c.NilNode(n)
diff --git a/conf/config.go b/conf/config.go
index e543732c..79989810 100644
--- a/conf/config.go
+++ b/conf/config.go
@@ -20,6 +20,7 @@ type Config struct {
 	ExpectAny   bool
 	Optimize    bool
 	Strict      bool
+	Profile     bool
 	ConstFns    map[string]reflect.Value
 	Visitors    []ast.Visitor
 	Functions   FunctionsTable
diff --git a/vm/opcodes.go b/vm/opcodes.go
index 0417dab6..84d751d6 100644
--- a/vm/opcodes.go
+++ b/vm/opcodes.go
@@ -81,6 +81,8 @@ const (
 	OpGroupBy
 	OpSortBy
 	OpSort
+	OpProfileStart
+	OpProfileEnd
 	OpBegin
 	OpEnd // This opcode must be at the end of this list.
 )
diff --git a/vm/program.go b/vm/program.go
index 4a878267..98954674 100644
--- a/vm/program.go
+++ b/vm/program.go
@@ -27,6 +27,7 @@ type Program struct {
 	variables int
 	functions []Function
 	debugInfo map[string]string
+	span      *Span
 }
 
 // NewProgram returns a new Program. It's used by the compiler.
@@ -40,6 +41,7 @@ func NewProgram(
 	arguments []int,
 	functions []Function,
 	debugInfo map[string]string,
+	span *Span,
 ) *Program {
 	return &Program{
 		source:    source,
@@ -51,6 +53,7 @@ func NewProgram(
 		Arguments: arguments,
 		functions: functions,
 		debugInfo: debugInfo,
+		span:      span,
 	}
 }
 
@@ -360,6 +363,12 @@ func (program *Program) DisassembleWriter(w io.Writer) {
 		case OpSort:
 			code("OpSort")
 
+		case OpProfileStart:
+			code("OpProfileStart")
+
+		case OpProfileEnd:
+			code("OpProfileEnd")
+
 		case OpBegin:
 			code("OpBegin")
 
diff --git a/vm/utils.go b/vm/utils.go
index d7db2a52..fc2f5e7b 100644
--- a/vm/utils.go
+++ b/vm/utils.go
@@ -2,6 +2,7 @@ package vm
 
 import (
 	"reflect"
+	"time"
 )
 
 type (
@@ -25,3 +26,15 @@ type Scope struct {
 }
 
 type groupBy = map[any][]any
+
+type Span struct {
+	Name       string  `json:"name"`
+	Expression string  `json:"expression"`
+	Duration   int64   `json:"duration"`
+	Children   []*Span `json:"children"`
+	start      time.Time
+}
+
+func GetSpan(program *Program) *Span {
+	return program.span
+}
diff --git a/vm/vm.go b/vm/vm.go
index 1e85893b..7e933ce7 100644
--- a/vm/vm.go
+++ b/vm/vm.go
@@ -8,6 +8,7 @@ import (
 	"regexp"
 	"sort"
 	"strings"
+	"time"
 
 	"github.com/expr-lang/expr/builtin"
 	"github.com/expr-lang/expr/file"
@@ -523,6 +524,14 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 			vm.memGrow(uint(scope.Len))
 			vm.push(sortable.Array)
 
+		case OpProfileStart:
+			span := program.Constants[arg].(*Span)
+			span.start = time.Now()
+
+		case OpProfileEnd:
+			span := program.Constants[arg].(*Span)
+			span.Duration += time.Since(span.start).Nanoseconds()
+
 		case OpBegin:
 			a := vm.pop()
 			array := reflect.ValueOf(a)