Add expr.Patch func

Author: antonmedv

ast/node.go

@@ -7,6 +7,12 @@ import (
 	"github.com/antonmedv/expr/file"
 )
 
+func Patch(node *Node, newNode Node) {
+	newNode.SetType((*node).Type())
+	newNode.SetLocation((*node).Location())
+	*node = newNode
+}
+
 // Node represents items of abstract syntax tree.
 type Node interface {
 	Location() file.Location

expr.go

@@ -2,6 +2,7 @@ package expr
 
 import (
 	"fmt"
+	"github.com/antonmedv/expr/ast"
 	"github.com/antonmedv/expr/file"
 	"reflect"
 
@@ -113,6 +114,13 @@ func Optimize(b bool) Option {
 	}
 }
 
+// TODO
+func Patch(visitor ast.Visitor) Option {
+	return func(c *conf.Config) {
+		c.Visitors = append(c.Visitors, visitor)
+	}
+}
+
 // Compile parses and compiles given input expression to bytecode program.
 func Compile(input string, ops ...Option) (*vm.Program, error) {
 	config := &conf.Config{
@@ -142,6 +150,10 @@ func Compile(input string, ops ...Option) (*vm.Program, error) {
 	// Patch operators before Optimize, as we may also mark it as ConstExpr.
 	compiler.PatchOperators(&tree.Node, config)
 
+	for _, v := range config.Visitors {
+		ast.Walk(&tree.Node, v)
+	}
+
 	if config.Optimize {
 		err = optimizer.Optimize(&tree.Node, config)
 		if err != nil {

expr_test.go

@@ -2,6 +2,7 @@ package expr_test
 
 import (
 	"fmt"
+	"github.com/antonmedv/expr/ast"
 	"strings"
 	"testing"
 	"time"
@@ -943,6 +944,53 @@ func TestConstExpr_error_no_env(t *testing.T) {
 	require.Equal(t, "no environment for const expression: divide", err.Error())
 }
 
+func TestPatch(t *testing.T) {
+	env := map[string]interface{}{
+		"x": map[string]interface{}{
+			"foo": map[string]interface{}{
+				"bar": "hello",
+			},
+		},
+	}
+
+	program, err := expr.Compile(
+		`x.foo.bar`,
+		expr.Env(env),
+		expr.Patch(&patcher{}),
+	)
+	require.NoError(t, err)
+
+	env = map[string]interface{}{
+		"x": map[string]interface{}{
+			"Foo": func() interface{} {
+				return map[string]interface{}{
+					"Bar": func() string {
+						return "patched"
+					},
+				}
+			},
+		},
+	}
+
+	out, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, "patched", out)
+}
+
+type patcher struct {
+}
+
+func (p *patcher) Enter(_ *ast.Node) {}
+func (p *patcher) Exit(node *ast.Node) {
+	switch n := (*node).(type) {
+	case *ast.PropertyNode:
+		ast.Patch(node, &ast.MethodNode{
+			Node:   n.Node,
+			Method: strings.Title(n.Property),
+		})
+	}
+}
+
 //
 // Mock types
 //

internal/conf/config.go

@@ -2,6 +2,7 @@ package conf
 
 import (
 	"fmt"
+	"github.com/antonmedv/expr/ast"
 	"github.com/antonmedv/expr/vm"
 	"reflect"
 )
@@ -16,6 +17,7 @@ type Config struct {
 	Strict       bool
 	DefaultType  reflect.Type
 	ConstExprFns map[string]reflect.Value
+	Visitors     []ast.Visitor
 	err          error
 }