Add Node method to Visitor interface to allow to patch ast tree

Author: davidklassen

ast/base_visitor.go

@@ -41,3 +41,5 @@ func (BaseVisitor) ArrayNode(node *ArrayNode) {}
 func (BaseVisitor) MapNode(node *MapNode) {}
 
 func (BaseVisitor) PairNode(node *PairNode) {}
+
+func (BaseVisitor) Node(node *Node) {}

ast/visitor.go

@@ -23,21 +23,24 @@ type Visitor interface {
 	ArrayNode(node *ArrayNode)
 	MapNode(node *MapNode)
 	PairNode(node *PairNode)
+	Node(node *Node)
 }
 
 type walker struct {
 	visitor Visitor
 }
 
-func Walk(node Node, visitor Visitor) {
+func Walk(node *Node, visitor Visitor) {
 	w := walker{
 		visitor: visitor,
 	}
 	w.walk(node)
 }
 
-func (w *walker) walk(node Node) {
-	switch n := node.(type) {
+func (w *walker) walk(node *Node) {
+	w.visitor.Node(node)
+
+	switch n := (*node).(type) {
 	case *NilNode:
 		w.visitor.NilNode(n)
 	case *IdentifierNode:
@@ -51,61 +54,62 @@ func (w *walker) walk(node Node) {
 	case *StringNode:
 		w.visitor.StringNode(n)
 	case *UnaryNode:
-		w.walk(n.Node)
+		w.walk(&n.Node)
 		w.visitor.UnaryNode(n)
 	case *BinaryNode:
-		w.walk(n.Left)
-		w.walk(n.Right)
+		w.walk(&n.Left)
+		w.walk(&n.Right)
 		w.visitor.BinaryNode(n)
 	case *MatchesNode:
-		w.walk(n.Left)
-		w.walk(n.Right)
+		w.walk(&n.Left)
+		w.walk(&n.Right)
 		w.visitor.MatchesNode(n)
 	case *PropertyNode:
-		w.walk(n.Node)
+		w.walk(&n.Node)
 		w.visitor.PropertyNode(n)
 	case *IndexNode:
-		w.walk(n.Node)
-		w.walk(n.Index)
+		w.walk(&n.Node)
+		w.walk(&n.Index)
 		w.visitor.IndexNode(n)
 	case *MethodNode:
-		w.walk(n.Node)
+		w.walk(&n.Node)
 		for _, arg := range n.Arguments {
-			w.walk(arg)
+			w.walk(&arg)
 		}
 		w.visitor.MethodNode(n)
 	case *FunctionNode:
 		for _, arg := range n.Arguments {
-			w.walk(arg)
+			w.walk(&arg)
 		}
 		w.visitor.FunctionNode(n)
 	case *BuiltinNode:
 		for _, arg := range n.Arguments {
-			w.walk(arg)
+			w.walk(&arg)
 		}
 		w.visitor.BuiltinNode(n)
 	case *ClosureNode:
-		w.walk(n.Node)
+		w.walk(&n.Node)
 		w.visitor.ClosureNode(n)
 	case *PointerNode:
 		w.visitor.PointerNode(n)
 	case *ConditionalNode:
-		w.walk(n.Cond)
-		w.walk(n.Exp1)
-		w.walk(n.Exp2)
+		w.walk(&n.Cond)
+		w.walk(&n.Exp1)
+		w.walk(&n.Exp2)
 		w.visitor.ConditionalNode(n)
 	case *ArrayNode:
 		for _, node := range n.Nodes {
-			w.walk(node)
+			w.walk(&node)
 		}
 		w.visitor.ArrayNode(n)
 	case *MapNode:
-		for _, pair := range n.Pairs {
-			w.walk(pair)
+		var pair Node
+		for _, pair = range n.Pairs {
+			w.walk(&pair)
 		}
 		w.visitor.MapNode(n)
 	case *PairNode:
-		w.walk(n.Value)
+		w.walk(&n.Value)
 		w.visitor.PairNode(n)
 	default:
 		panic(fmt.Sprintf("undefined node type (%T)", node))

ast/visitor_test.go

@@ -16,13 +16,38 @@ func (v *visitor) IdentifierNode(node *ast.IdentifierNode) {
 }
 
 func TestWalk(t *testing.T) {
-	node := &ast.BinaryNode{
+	var node ast.Node
+	node = &ast.BinaryNode{
 		Operator: "+",
 		Left:     &ast.IdentifierNode{Value: "foo"},
 		Right:    &ast.IdentifierNode{Value: "bar"},
 	}
 
 	visitor := &visitor{}
-	ast.Walk(node, visitor)
+	ast.Walk(&node, visitor)
 	assert.Equal(t, []string{"foo", "bar"}, visitor.identifiers)
 }
+
+type patcher struct {
+	ast.BaseVisitor
+}
+
+func (p *patcher) Node(node *ast.Node) {
+	if _, ok := (*node).(*ast.IdentifierNode); ok {
+		*node = &ast.NilNode{}
+	}
+}
+
+func TestWalk_patch(t *testing.T) {
+	var node ast.Node
+	node = &ast.BinaryNode{
+		Operator: "+",
+		Left:     &ast.IdentifierNode{Value: "foo"},
+		Right:    &ast.IdentifierNode{Value: "bar"},
+	}
+
+	patcher := &patcher{}
+	ast.Walk(&node, patcher)
+	assert.IsType(t, &ast.NilNode{}, node.(*ast.BinaryNode).Left)
+	assert.IsType(t, &ast.NilNode{}, node.(*ast.BinaryNode).Right)
+}

checker/types.go

@@ -281,7 +281,7 @@ func (v setVisitor) IntegerNode(node *ast.IntegerNode) {
 }
 
 func setUncertainType(node ast.Node, t reflect.Type) {
-	ast.Walk(node, setVisitor{t: dereference(t)})
+	ast.Walk(&node, setVisitor{t: dereference(t)})
 }
 
 type hasVisitor struct {
@@ -295,6 +295,6 @@ func (v *hasVisitor) IntegerNode(node *ast.IntegerNode) {
 
 func isCertain(node ast.Node) bool {
 	v := &hasVisitor{certain: true}
-	ast.Walk(node, v)
+	ast.Walk(&node, v)
 	return v.certain
 }

cmd/ecc/dot.go

@@ -18,12 +18,13 @@ const format = `digraph {
 
 func dotAst(node Node) {
 	v := &visitor{}
-	Walk(node, v)
+	Walk(&node, v)
 	dot := fmt.Sprintf(format, v.nodes, v.links)
 	_, _ = fmt.Fprintf(os.Stdout, dot)
 }
 
 type visitor struct {
+	BaseVisitor
 	nodes string
 	links string
 	index int