Add ast.Find

Author: antonmedv

ast/find.go

@@ -0,0 +1,18 @@
+package ast
+
+func Find(node Node, fn func(node Node) bool) Node {
+	v := &finder{fn: fn}
+	Walk(&node, v)
+	return v.node
+}
+
+type finder struct {
+	node Node
+	fn   func(node Node) bool
+}
+
+func (f *finder) Visit(node *Node) {
+	if f.fn(*node) {
+		f.node = *node
+	}
+}

ast/find_test.go

@@ -0,0 +1,31 @@
+package ast_test
+
+import (
+	"testing"
+
+	"github.com/expr-lang/expr/internal/testify/require"
+
+	"github.com/expr-lang/expr/ast"
+)
+
+func TestFind(t *testing.T) {
+	left := &ast.IdentifierNode{
+		Value: "a",
+	}
+	var root ast.Node = &ast.BinaryNode{
+		Operator: "+",
+		Left:     left,
+		Right: &ast.IdentifierNode{
+			Value: "b",
+		},
+	}
+
+	x := ast.Find(root, func(node ast.Node) bool {
+		if n, ok := node.(*ast.IdentifierNode); ok {
+			return n.Value == "a"
+		}
+		return false
+	})
+
+	require.Equal(t, left, x)
+}