Repository: openllb/hlb
Pull request #302 (2022-02-14 20:43:24)

Refactor ast.Find to ast.Search, introduce new ast.Find and tests
- Refactors `ast.Find` to `ast.Search`
- Introduces new `ast.Find` that finds by line + column rather than by `strings.Contains`
- Adds tests for `ast.Match`, `ast.Search`, `ast.Find`, and `ast.IsPositionWithinNode`


Patch:
diff --git a/langserver/server.go b/langserver/server.go
index 35233ad..5db2564 100644
--- a/langserver/server.go
+++ b/langserver/server.go
@@ -707,13 +707,7 @@ func (ls *LangServer) textDocumentCompletionHandler(ctx context.Context, params
 }
 
 func isPositionWithinNode(pos lsp.Position, node ast.Node) bool {
-	if (pos.Line < node.Position().Line-1 || pos.Line > node.End().Line-1) ||
-		(pos.Line == node.Position().Line-1 && pos.Character < node.Position().Column-1) ||
-		(pos.Line == node.End().Line-1 && pos.Character >= node.End().Column-1) {
-		return false
-	}
-
-	return true
+	return ast.IsPositionWithinNode(node, pos.Line-1, pos.Character-1)
 }
 
 func newLocationFromNode(uri lsp.DocumentURI, node ast.Node) *lsp.Location {
diff --git a/parser/ast/context.go b/parser/ast/context.go
deleted file mode 100644
index 28936d3..0000000
--- a/parser/ast/context.go
+++ /dev/null
@@ -1,43 +0,0 @@
-package ast
-
-import (
-	"context"
-	"sync"
-)
-
-type modulesKey struct{}
-
-func WithModules(ctx context.Context, mods *ModuleLookup) context.Context {
-	return context.WithValue(ctx, modulesKey{}, mods)
-}
-
-func Modules(ctx context.Context) *ModuleLookup {
-	mods, ok := ctx.Value(modulesKey{}).(*ModuleLookup)
-	if !ok {
-		return NewModules()
-	}
-	return mods
-}
-
-type ModuleLookup struct {
-	mods map[string]*Module
-	mu   sync.RWMutex
-}
-
-func NewModules() *ModuleLookup {
-	return &ModuleLookup{
-		mods: make(map[string]*Module),
-	}
-}
-
-func (ml *ModuleLookup) Get(filename string) *Module {
-	ml.mu.RLock()
-	defer ml.mu.RUnlock()
-	return ml.mods[filename]
-}
-
-func (ml *ModuleLookup) Set(filename string, mod *Module) {
-	ml.mu.Lock()
-	defer ml.mu.Unlock()
-	ml.mods[filename] = mod
-}
diff --git a/parser/ast/match.go b/parser/ast/match.go
new file mode 100644
index 0000000..0c5ca1b
--- /dev/null
+++ b/parser/ast/match.go
@@ -0,0 +1,145 @@
+package ast
+
+import (
+	"fmt"
+	"reflect"
+)
+
+// MatchOpts provides options while walking the CST.
+type MatchOpts struct {
+	// Filter is called to see if the node should be walked. If nil, then all
+	// nodes are walked.
+	Filter func(Node) bool
+
+	// AllowDuplicates is enabled to allow duplicate CST structures between
+	// arguments in the functions provided.
+	//
+	// For example, if a function is defined as `func(*FuncDecl, *CallStmt)`
+	// sequences like `[*FuncDecl, ..., *CallStmt, ..., *CallStmt]` are not
+	// matched by default. Allowing duplicates will match instead.
+	AllowDuplicates bool
+}
+
+type matcher struct {
+	opts    MatchOpts
+	vs      []reflect.Value
+	expects [][]reflect.Type
+	actuals [][]reflect.Value
+	indices []int
+}
+
+// Match walks a CST and invokes given functions if their arguments match a
+// non-contiguous sequence of current path walked. This is useful when you want
+// to walk to a specific type of Node, while having access to specific parents
+// of the Node.
+//
+// The function arguments must all implement the Node interface, and may be
+// a non-contiguous sequence. That is, you don't have to specify every CST
+// structure.
+//
+// The sequence is matched right to left, from the deepest node first. The
+// final argument will always be the current node being visited.
+//
+// When multiple functions are matched, they are invoked in the order given
+// to Match. That way, you can write functions that positively match, and then
+// provide a more general function as a catch all without walking the CST a
+// second time.
+//
+// For example, you can invoke Match to find CallStmts inside FuncLits:
+// ```go
+// Match(root, MatchOpts{},
+// 	func(lit *FuncLit, call *CallStmt) {
+// 		fmt.Println(lit.Pos, call.Pos)
+// 	},
+// )
+// ```
+func Match(root Node, opts MatchOpts, funs ...interface{}) {
+	m := &matcher{
+		opts:    opts,
+		vs:      make([]reflect.Value, len(funs)),
+		expects: make([][]reflect.Type, len(funs)),
+		actuals: make([][]reflect.Value, len(funs)),
+		indices: make([]int, len(funs)),
+	}
+
+	for i, fun := range funs {
+		m.vs[i] = reflect.ValueOf(fun)
+	}
+
+	node := reflect.TypeOf((*Node)(nil)).Elem()
+	for i, v := range m.vs {
+		t := v.Type()
+		for j := 0; j < t.NumIn(); j++ {
+			arg := t.In(j)
+			if !arg.Implements(node) {
+				panic(fmt.Sprintf("%s has bad signature: %s does not implement Node", t, arg))
+			}
+
+			m.expects[i] = append(m.expects[i], arg)
+		}
+	}
+
+	for i, expect := range m.expects {
+		m.actuals[i] = make([]reflect.Value, len(expect))
+	}
+
+	Walk(root, m)
+}
+
+func (m *matcher) Visit(in Introspector, n Node) Visitor {
+	if n == nil {
+		return nil
+	}
+
+	if m.opts.Filter != nil {
+		if !m.opts.Filter(n) {
+			return nil
+		}
+	}
+
+	// Clear out indices from a previous visit.
+	for i := 0; i < len(m.expects); i++ {
+		m.indices[i] = len(m.expects[i]) - 1
+	}
+
+	for i := len(in.Path()) - 1; i >= 0; i-- {
+		p := in.Path()[i]
+		v := reflect.ValueOf(p)
+
+		for j, expect := range m.expects {
+			k := m.indices[j]
+
+			// Either the function has been matched or will never match.
+			if k < 0 {
+				continue
+			}
+
+			if v.Type() != expect[k] {
+				if i == len(in.Path())-1 {
+					// The final argument must always match the deepest node.
+					m.indices[j] = -2
+				} else if !m.opts.AllowDuplicates && v.Type() == expect[k+1] {
+					// Unless duplicates are allowed, the current node cannot be the same
+					// type as the previous matched node.
+					m.indices[j] = -2
+				}
+
+				continue
+			}
+
+			m.actuals[j][k] = v
+			m.indices[j] -= 1
+		}
+	}
+
+	// Invoke matched functions in the order they were given.
+	for i := 0; i < len(m.vs); i++ {
+		// Functions that will never match have an index of -2.
+		// Functions that matched have an index of -1.
+		if m.indices[i] == -1 {
+			m.vs[i].Call(m.actuals[i])
+		}
+	}
+
+	return m
+}
diff --git a/parser/ast/walk.go b/parser/ast/walk.go
index 917502c..6c58905 100644
--- a/parser/ast/walk.go
+++ b/parser/ast/walk.go
@@ -1,9 +1,9 @@
 package ast
 
 import (
-	"fmt"
-	"reflect"
 	"strings"
+
+	"github.com/alecthomas/participle/lexer"
 )
 
 // A Visitor's Visit method is invoked for each node encountered by Walk.
@@ -336,183 +336,116 @@ func (w *walker) walkHeredocFragments(list []*HeredocFragment, v Visitor) {
 	}
 }
 
-// MatchOpts provides options while walking the CST.
-type MatchOpts struct {
-	// Filter is called to see if the node should be walked. If nil, then all
-	// nodes are walked.
-	Filter func(Node) bool
-
-	// AllowDuplicates is enabled to allow duplicate CST structures between
-	// arguments in the functions provided.
-	//
-	// For example, if a function is defined as `func(*FuncDecl, *CallStmt)`
-	// sequences like `[*FuncDecl, ..., *CallStmt, ..., *CallStmt]` are not
-	// matched by default. Allowing duplicates will match instead.
-	AllowDuplicates bool
-}
-
-type matcher struct {
-	opts    MatchOpts
-	vs      []reflect.Value
-	expects [][]reflect.Type
-	actuals [][]reflect.Value
-	indices []int
-}
-
-// Match walks a CST and invokes given functions if their arguments match a
-// non-contiguous sequence of current path walked. This is useful when you want
-// to walk to a specific type of Node, while having access to specific parents
-// of the Node.
-//
-// The function arguments must all implement the Node interface, and may be
-// a non-contiguous sequence. That is, you don't have to specify every CST
-// structure.
-//
-// The sequence is matched right to left, from the deepest node first. The
-// final argument will always be the current node being visited.
-//
-// When multiple functions are matched, they are invoked in the order given
-// to Match. That way, you can write functions that positively match, and then
-// provide a more general function as a catch all without walking the CST a
-// second time.
-//
-// For example, you can invoke Match to find CallStmts inside FuncLits:
-// ```go
-// Match(root, MatchOpts{},
-// 	func(lit *FuncLit, call *CallStmt) {
-// 		fmt.Println(lit.Pos, call.Pos)
-// 	},
-// )
-// ```
-func Match(root Node, opts MatchOpts, fs ...interface{}) {
-	m := &matcher{
-		opts:    opts,
-		vs:      make([]reflect.Value, len(fs)),
-		expects: make([][]reflect.Type, len(fs)),
-		actuals: make([][]reflect.Value, len(fs)),
-		indices: make([]int, len(fs)),
-	}
-
-	for i, f := range fs {
-		m.vs[i] = reflect.ValueOf(f)
-	}
-
-	node := reflect.TypeOf((*Node)(nil)).Elem()
-	for i, v := range m.vs {
-		t := v.Type()
-		for j := 0; j < t.NumIn(); j++ {
-			arg := t.In(j)
-			if !arg.Implements(node) {
-				panic(fmt.Sprintf("%s has bad signature: %s does not implement Node", t, arg))
-			}
-
-			m.expects[i] = append(m.expects[i], arg)
-		}
-	}
-
-	for i, expect := range m.expects {
-		m.actuals[i] = make([]reflect.Value, len(expect))
-	}
-
-	Walk(root, m)
+type searcher struct {
+	match Node
+	query string
+	skip  int
 }
 
-func (m *matcher) Visit(in Introspector, n Node) Visitor {
+func (s *searcher) Visit(_ Introspector, n Node) Visitor {
 	if n == nil {
 		return nil
 	}
-
-	if m.opts.Filter != nil {
-		if !m.opts.Filter(n) {
+	if strings.Contains(n.String(), s.query) && s.skip >= 0 {
+		s.match = n
+		if n.String() == s.query {
+			s.skip -= 1
+			if s.skip >= 0 {
+				s.match = nil
+			}
 			return nil
+		} else {
+			return s
 		}
 	}
+	return nil
+}
 
-	// Clear out indices from a previous visit.
-	for i := 0; i < len(m.expects); i++ {
-		m.indices[i] = len(m.expects[i]) - 1
-	}
-
-	for i := len(in.Path()) - 1; i >= 0; i-- {
-		p := in.Path()[i]
-		v := reflect.ValueOf(p)
-
-		for j, expect := range m.expects {
-			k := m.indices[j]
-
-			// Either the function has been matched or will never match.
-			if k < 0 {
-				continue
-			}
-
-			if v.Type() != expect[k] {
-				if i == len(in.Path())-1 {
-					// The final argument must always match the deepest node.
-					m.indices[j] = -2
-				} else if !m.opts.AllowDuplicates && v.Type() == expect[k+1] {
-					// Unless duplicates are allowed, the current node cannot be the same
-					// type as the previous matched node.
-					m.indices[j] = -2
-				}
-
-				continue
-			}
+// SearchOption provides configuration for Search.
+type SearchOption func(*searcher)
 
-			m.actuals[j][k] = v
-			m.indices[j] -= 1
-		}
+// WithSkip specifies how many matches to skip before returning.
+func WithSkip(skip int) SearchOption {
+	return func(f *searcher) {
+		f.skip = skip
 	}
+}
 
-	// Invoke matched functions in the order they were given.
-	for i := 0; i < len(m.vs); i++ {
-		// Functions that will never match have an index of -2.
-		// Functions that matched have an index of -1.
-		if m.indices[i] == -1 {
-			m.vs[i].Call(m.actuals[i])
-		}
+// Search searches for the deepest node that contains the query.
+func Search(root Node, query string, opts ...SearchOption) Node {
+	f := &searcher{query: query}
+	for _, opt := range opts {
+		opt(f)
 	}
-
-	return m
+	Walk(root, f)
+	return f.match
 }
 
 type finder struct {
-	node  Node
-	match string
-	skip  int
+	match  Node
+	line   int
+	column int
+	filter func(Node) bool
 }
 
-func (v *finder) Visit(_ Introspector, n Node) Visitor {
+func (f *finder) Visit(_ Introspector, n Node) Visitor {
 	if n == nil {
 		return nil
 	}
-	if strings.Contains(n.String(), v.match) && v.skip >= 0 {
-		v.node = n
-		if n.String() == v.match {
-			v.skip -= 1
-			if v.skip >= 0 {
-				v.node = nil
+	if f.column == 0 {
+		if (f.line >= n.Position().Line && f.line <= n.End().Line) &&
+			(f.match == nil || n.End().Column <= f.match.End().Column) {
+			if f.line == n.Position().Line && f.line == n.End().Line {
+				keep := true
+				if f.filter != nil {
+					keep = f.filter(n)
+				}
+				if keep {
+					f.match = n
+				}
 			}
-			return nil
-		} else {
-			return v
+			return f
+		}
+	} else if IsPositionWithinNode(n, f.line, f.column) {
+		keep := true
+		if f.filter != nil {
+			keep = f.filter(n)
 		}
+		if keep {
+			f.match = n
+		}
+		return f
 	}
 	return nil
 }
 
-type FindOption func(*finder)
+// Find finds the deepest node that is on a specified line or column.
+//
+// If column is 0, it ignores column until it finds a match, applying the
+// matched node's start to end column constraints to its deeper matches.
+//
+// If filter is specified, each match will only be kept if filter returns true.
+func Find(root Node, line, column int, filter func(Node) bool) Node {
+	f := &finder{line: line, column: column, filter: filter}
+	Walk(root, f)
+	return f.match
+}
 
-func WithSkip(skip int) FindOption {
-	return func(f *finder) {
-		f.skip = skip
-	}
+// IsPositionWithinNode returns true if a line column is within a node's position.
+func IsPositionWithinNode(node Node, line, column int) bool {
+	return IsIntersect(node.Position(), node.End(), line, column)
 }
 
-func Find(root Node, match string, opts ...FindOption) Node {
-	f := &finder{match: match}
-	for _, opt := range opts {
-		opt(f)
+// IsIntersect returns true if a line column is within a start and end
+// position.
+func IsIntersect(start, end lexer.Position, line, column int) bool {
+	if start.Column == 0 || end.Column == 0 || column == 0 {
+		return line >= start.Line && line <= end.Line
 	}
-	Walk(root, f)
-	return f.node
+	if (line < start.Line || line > end.Line) ||
+		(line == start.Line && column < start.Column) ||
+		(line == end.Line && column >= end.Column) {
+		return false
+	}
+	return true
 }
diff --git a/parser/parse.go b/parser/parse.go
index 4db42aa..0fc6223 100644
--- a/parser/parse.go
+++ b/parser/parse.go
@@ -22,20 +22,22 @@ func Parse(ctx context.Context, r io.Reader) (*ast.Module, error) {
 	mod := &ast.Module{}
 	defer AssignDocStrings(mod)
 
-	fb := filebuffer.New(name)
-	r = io.TeeReader(r, fb)
-	defer func() {
-		if mod.Pos.Filename != "" {
-			filebuffer.Buffers(ctx).Set(mod.Pos.Filename, fb)
-		}
-	}()
+	files := filebuffer.Buffers(ctx)
+	if files != nil {
+		fb := filebuffer.New(name)
+		r = io.TeeReader(r, fb)
+		defer func() {
+			if mod.Pos.Filename != "" {
+				files.Set(mod.Pos.Filename, fb)
+			}
+		}()
+	}
 
 	err := ast.Parser.Parse(name, r, mod)
 	if err != nil {
 		return nil, err
 	}
 	mod.Directory = NewLocalDirectory(filepath.Dir(mod.Pos.Filename), "")
-	ast.Modules(ctx).Set(mod.Pos.Filename, mod)
 	return mod, nil
 }
 


Test patch:
diff --git a/checker/checker_test.go b/checker/checker_test.go
index ed047a0..661d482 100644
--- a/checker/checker_test.go
+++ b/checker/checker_test.go
@@ -171,8 +171,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithNumArgs(
-				ast.Find(mod, "image"), 1, 0,
-				errdefs.Defined(ast.Find(builtin.Module, "image")),
+				ast.Search(mod, "image"), 1, 0,
+				errdefs.Defined(ast.Search(builtin.Module, "image")),
 			)
 		},
 	}, {
@@ -185,8 +185,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDuplicates([]ast.Node{
-				ast.Find(mod, "duplicate"),
-				ast.Find(mod, "duplicate", ast.WithSkip(1)),
+				ast.Search(mod, "duplicate"),
+				ast.Search(mod, "duplicate", ast.WithSkip(1)),
 			})
 		},
 	}, {
@@ -201,8 +201,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDuplicates([]ast.Node{
-				ast.Find(mod, "duplicate"),
-				ast.Find(mod, "duplicate", ast.WithSkip(1)),
+				ast.Search(mod, "duplicate"),
+				ast.Search(mod, "duplicate", ast.WithSkip(1)),
 			})
 		},
 	}, {
@@ -214,8 +214,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDuplicates([]ast.Node{
-				ast.Find(builtin.Module, "image"),
-				ast.Find(mod, "image"),
+				ast.Search(builtin.Module, "image"),
+				ast.Search(mod, "image"),
 			})
 		},
 	}, {
@@ -229,8 +229,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDuplicates([]ast.Node{
-				ast.Find(builtin.Module, "image"),
-				ast.Find(mod, "image"),
+				ast.Search(builtin.Module, "image"),
+				ast.Search(mod, "image"),
 			})
 		},
 	}, {
@@ -247,8 +247,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDuplicates([]ast.Node{
-				ast.Find(mod, "duplicate"),
-				ast.Find(mod, "duplicate", ast.WithSkip(1)),
+				ast.Search(mod, "duplicate"),
+				ast.Search(mod, "duplicate", ast.WithSkip(1)),
 			})
 		},
 	}, {
@@ -262,8 +262,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithCallImport(
-				ast.Find(mod, "foo", ast.WithSkip(1)),
-				ast.Find(mod, "foo"),
+				ast.Search(mod, "foo", ast.WithSkip(1)),
+				ast.Search(mod, "foo"),
 			)
 		},
 	}, {
@@ -293,7 +293,7 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithUndefinedIdent(
-				ast.Find(mod, "foo"),
+				ast.Search(mod, "foo"),
 				nil,
 			)
 		},
@@ -307,8 +307,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithNotImport(
-				ast.Find(mod, "myFunction.build").(*ast.IdentExpr),
-				ast.Find(mod, "myFunction"),
+				ast.Search(mod, "myFunction.build").(*ast.IdentExpr),
+				ast.Search(mod, "myFunction"),
 			)
 		},
 	}, {
@@ -335,10 +335,10 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithWrongType(
-				ast.Find(mod, "image"),
+				ast.Search(mod, "image"),
 				[]ast.Kind{ast.Pipeline},
 				ast.Filesystem,
-				errdefs.Defined(ast.Find(builtin.Module, "image")),
+				errdefs.Defined(ast.Search(builtin.Module, "image")),
 			)
 		},
 	}, {
@@ -353,7 +353,7 @@ func TestChecker_Check(t *testing.T) {
 		}
 		`,
 		func(mod *ast.Module) error {
-			return errdefs.WithNoBindTarget(ast.Find(mod, "as"))
+			return errdefs.WithNoBindTarget(ast.Search(mod, "as"))
 		},
 	}, {
 		"no error when bind list is empty",
@@ -373,10 +373,10 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithWrongType(
-				ast.Find(mod, "imageID", ast.WithSkip(1)),
+				ast.Search(mod, "imageID", ast.WithSkip(1)),
 				[]ast.Kind{ast.Filesystem},
 				ast.String,
-				errdefs.Defined(ast.Find(mod, "imageID")),
+				errdefs.Defined(ast.Search(mod, "imageID")),
 			)
 		},
 	}, {
@@ -389,10 +389,10 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithWrongType(
-				ast.Find(mod, "imageID", ast.WithSkip(1)),
+				ast.Search(mod, "imageID", ast.WithSkip(1)),
 				[]ast.Kind{ast.Filesystem},
 				ast.String,
-				errdefs.Defined(ast.Find(mod, "imageID")),
+				errdefs.Defined(ast.Search(mod, "imageID")),
 			)
 		},
 	}, {
@@ -404,10 +404,10 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithNoBindEffects(
-				ast.Find(mod, "run"),
-				ast.Find(mod, "as"),
+				ast.Search(mod, "run"),
+				ast.Search(mod, "as"),
 				errdefs.Defined(
-					ast.Find(builtin.Module, "run"),
+					ast.Search(builtin.Module, "run"),
 				),
 			)
 		},
@@ -420,8 +420,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithUndefinedBindTarget(
-				ast.Find(mod, "dockerPush"),
-				ast.Find(mod, "undefined"),
+				ast.Search(mod, "dockerPush"),
+				ast.Search(mod, "undefined"),
 			)
 		},
 	}, {
@@ -433,8 +433,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithNoBindClosure(
-				ast.Find(mod, "as"),
-				ast.Find(mod, "option::run"),
+				ast.Search(mod, "as"),
+				ast.Search(mod, "option::run"),
 			)
 		},
 	}, {
@@ -452,8 +452,8 @@ func TestChecker_Check(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithNoBindClosure(
-				ast.Find(mod, "as"),
-				ast.Find(mod, "option::run"),
+				ast.Search(mod, "as"),
+				ast.Search(mod, "option::run"),
 			)
 		},
 	}, {
diff --git a/codegen/codegen_test.go b/codegen/codegen_test.go
index 23efd79..eab7470 100644
--- a/codegen/codegen_test.go
+++ b/codegen/codegen_test.go
@@ -1022,7 +1022,7 @@ func TestCodegenError(t *testing.T) {
 			func(mod *ast.Module) error {
 				return errdefs.WithInvalidImageRef(
 					errors.New("invalid reference format"),
-					ast.Find(mod, `"#"`),
+					ast.Search(mod, `"#"`),
 					"#",
 				)
 			},
@@ -1109,9 +1109,9 @@ func TestCodeGenImport(t *testing.T) {
 		}},
 		func(mod *ast.Module) error {
 			return errdefs.WithUndefinedIdent(
-				ast.Find(mod, "bar"),
+				ast.Search(mod, "bar"),
 				nil,
-				errdefs.Imported(ast.Find(mod, "other")),
+				errdefs.Imported(ast.Search(mod, "other")),
 			)
 		},
 	}, {
@@ -1131,8 +1131,8 @@ func TestCodeGenImport(t *testing.T) {
 		}},
 		func(mod *ast.Module) error {
 			return errdefs.WithCallUnexported(
-				ast.Find(mod, "foo"),
-				errdefs.Imported(ast.Find(mod, "other")),
+				ast.Search(mod, "foo"),
+				errdefs.Imported(ast.Search(mod, "other")),
 			)
 		},
 	}, {
diff --git a/linter/linter_test.go b/linter/linter_test.go
index 9562c4e..ce63fc9 100644
--- a/linter/linter_test.go
+++ b/linter/linter_test.go
@@ -32,7 +32,7 @@ func TestLinter_Lint(t *testing.T) {
 		`,
 		func(mod *ast.Module) error {
 			return errdefs.WithDeprecated(
-				mod, ast.Find(mod, `"./foo.hlb"`).(*ast.StringLit),
+				mod, ast.Search(mod, `"./foo.hlb"`).(*ast.StringLit),
 				`import path without keyword "from" is deprecated`,
 			)
 		},
@@ -50,11 +50,11 @@ func TestLinter_Lint(t *testing.T) {
 			return &diagnostic.Error{
 				Diagnostics: []error{
 					errdefs.WithDeprecated(
-						mod, ast.Find(mod, "group"),
+						mod, ast.Search(mod, "group"),
 						"type `group` is deprecated, use `pipeline` instead",
 					),
 					errdefs.WithDeprecated(
-						mod, ast.Find(mod, "parallel"),
+						mod, ast.Search(mod, "parallel"),
 						"function `parallel` is deprecated, use `stage` instead",
 					),
 				},
diff --git a/module/resolve_test.go b/module/resolve_test.go
index 1ed7062..a46121f 100644
--- a/module/resolve_test.go
+++ b/module/resolve_test.go
@@ -95,7 +95,7 @@ func TestResolveGraph(t *testing.T) {
 		func(mod *ast.Module, imods map[string]*ast.Module) error {
 			return errdefs.WithImportPathNotExist(
 				os.ErrNotExist,
-				ast.Find(mod, `"unknown.hlb"`),
+				ast.Search(mod, `"unknown.hlb"`),
 				"unknown.hlb",
 			)
 		},
@@ -107,7 +107,7 @@ func TestResolveGraph(t *testing.T) {
 		func(mod *ast.Module, imods map[string]*ast.Module) error {
 			return errdefs.WithImportPathNotExist(
 				os.ErrNotExist,
-				ast.Find(imods["transitive"], `"unknown.hlb"`),
+				ast.Search(imods["transitive"], `"unknown.hlb"`),
 				"unknown.hlb",
 			)
 		},
diff --git a/parser/ast/match_test.go b/parser/ast/match_test.go
new file mode 100644
index 0000000..51c00ec
--- /dev/null
+++ b/parser/ast/match_test.go
@@ -0,0 +1,188 @@
+package ast
+
+import (
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type matches []match
+
+type match []string
+
+func matched(ifaces ...interface{}) match {
+	var types []string
+	for _, iface := range ifaces {
+		types = append(types, reflect.TypeOf(iface).String())
+	}
+	return match(types)
+}
+
+func TestMatch(t *testing.T) {
+	type testCase struct {
+		name     string
+		input    string
+		fun      func(Node, chan match)
+		expected matches
+	}
+
+	for _, tc := range []testCase{{
+		"empty",
+		``,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{})
+		},
+		nil,
+	}, {
+		"root matcher",
+		``,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{},
+				func(mod *Module) {
+					ms <- matched(mod)
+				},
+			)
+		},
+		matches{
+			matched(&Module{}),
+		},
+	}, {
+		"single matcher",
+		`
+		fs foo()
+		fs bar()
+		`,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{},
+				func(fd *FuncDecl) {
+					ms <- matched(fd)
+				},
+			)
+		},
+		matches{
+			matched(&FuncDecl{}),
+			matched(&FuncDecl{}),
+		},
+	}, {
+		"chain matcher",
+		`
+		fs default() {
+			image "alpine"
+			run "echo foo" with option {
+				env "key" "value"
+			}
+		}
+		`,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{},
+				func(parentCall *CallStmt, childCall *CallStmt) {
+					ms <- matched(parentCall, childCall)
+				},
+			)
+		},
+		matches{
+			matched(&CallStmt{}, &CallStmt{}),
+		},
+	}, {
+		"chain matcher with nodes between",
+		`
+		fs default() {
+			image "alpine"
+			run "echo foo" with option {
+				env "key" "value"
+			}
+		}
+		`,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{},
+				func(fd *FuncDecl, call *CallStmt) {
+					ms <- matched(fd, call)
+				},
+			)
+		},
+		matches{
+			// image "alpine"
+			matched(&FuncDecl{}, &CallStmt{}),
+			// run "echo foo"
+			matched(&FuncDecl{}, &CallStmt{}),
+		},
+	}, {
+		"chain matcher with nodes between but allow duplicates",
+		`
+		fs default() {
+			image "alpine"
+			run "echo foo" with option {
+				env "key" "value"
+			}
+		}
+		`,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{
+				AllowDuplicates: true,
+			}, func(fd *FuncDecl, call *CallStmt) {
+				ms <- matched(fd, call)
+			},
+			)
+		},
+		matches{
+			// image "alpine"
+			matched(&FuncDecl{}, &CallStmt{}),
+			// run "echo foo"
+			matched(&FuncDecl{}, &CallStmt{}),
+			// env "key" "value"
+			matched(&FuncDecl{}, &CallStmt{}),
+		},
+	}, {
+		"multiple matchers",
+		`
+		import util from fs {
+			image "util.hlb"
+		}
+
+		fs default() {
+			util.base
+			run "echo foo" with option {
+				env "key" "value"
+			}
+		}
+		`,
+		func(root Node, ms chan match) {
+			Match(root, MatchOpts{},
+				func(id *ImportDecl, lit *FuncLit) {
+					ms <- matched(id, lit)
+				},
+				func(fd *FuncDecl, lit *FuncLit) {
+					ms <- matched(fd, lit)
+				},
+			)
+		},
+		matches{
+			// from fs { ... }
+			matched(&ImportDecl{}, &FuncLit{}),
+			// with option { ... }
+			matched(&FuncDecl{}, &FuncLit{}),
+		},
+	}} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			mod := &Module{}
+			r := strings.NewReader(cleanup(tc.input))
+			err := Parser.Parse("", r, mod)
+			require.NoError(t, err)
+
+			ms := make(chan match)
+			go func() {
+				defer close(ms)
+				tc.fun(mod, ms)
+			}()
+
+			var actual matches
+			for m := range ms {
+				actual = append(actual, m)
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
diff --git a/parser/ast/walk_test.go b/parser/ast/walk_test.go
new file mode 100644
index 0000000..b5f2faa
--- /dev/null
+++ b/parser/ast/walk_test.go
@@ -0,0 +1,236 @@
+package ast
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/alecthomas/participle/lexer"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSearch(t *testing.T) {
+	type testCase struct {
+		name     string
+		input    string
+		query    string
+		opts     []SearchOption
+		expected string
+	}
+
+	for _, tc := range []testCase{{
+		"empty",
+		``,
+		"",
+		nil,
+		"1:1",
+	}, {
+		"specific match",
+		`
+		fs foo()
+		`,
+		"foo",
+		nil,
+		"1:4",
+	}, {
+		"partial match",
+		`
+		fs foo()
+		`,
+		"foo()",
+		nil,
+		"1:1",
+	}, {
+		"skip match",
+		`
+		fs default() {
+			image "alpine"
+			run "echo hello"
+			run "echo world"
+		}
+		`,
+		`run`,
+		[]SearchOption{WithSkip(1)},
+		"4:2",
+	}} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			mod := &Module{}
+			r := strings.NewReader(cleanup(tc.input))
+			err := Parser.Parse("", r, mod)
+			require.NoError(t, err)
+
+			node := Search(mod, tc.query, tc.opts...)
+			actual := ""
+			if node != nil {
+				actual = formatPos(node.Position())
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
+func TestFind(t *testing.T) {
+	type testCase struct {
+		name                       string
+		input                      string
+		line, column               int
+		filter                     func(Node) bool
+		expectedStart, expectedEnd string
+	}
+
+	for _, tc := range []testCase{{
+		"empty",
+		``,
+		0, 0,
+		nil,
+		"", "",
+	}, {
+		"no column match",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		2, 0,
+		nil,
+		"2:2", "2:7", // image "alpine"
+	}, {
+		"column match node",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		2, 3,
+		nil,
+		"2:2", "2:7", // image "alpine"
+	}, {
+		"column match parent",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		2, 1,
+		nil,
+		"1:14", "3:2", // parent block stmt { ... }
+	}, {
+		"filter match",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		1, 0,
+		func(n Node) bool {
+			_, ok := n.(*FuncSignature)
+			return ok
+		},
+		"1:1", "1:13", // fs default()
+	}} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			mod := &Module{}
+			r := strings.NewReader(cleanup(tc.input))
+			err := Parser.Parse("", r, mod)
+			require.NoError(t, err)
+
+			node := Find(mod, tc.line, tc.column, tc.filter)
+			actualStart := ""
+			actualEnd := ""
+			if node != nil {
+				actualStart = formatPos(node.Position())
+				actualEnd = formatPos(node.End())
+			}
+			require.Equal(t, tc.expectedStart, actualStart)
+			require.Equal(t, tc.expectedEnd, actualEnd)
+		})
+	}
+}
+
+func formatPos(pos lexer.Position) string {
+	return fmt.Sprintf("%d:%d", pos.Line, pos.Column)
+}
+
+func TestIsPositionWithinNode(t *testing.T) {
+	type testCase struct {
+		name         string
+		input        string
+		match        func(root Node) Node
+		line, column int
+		expected     bool
+	}
+
+	for _, tc := range []testCase{{
+		"empty",
+		``,
+		func(root Node) Node {
+			return root
+		},
+		1, 1,
+		true,
+	}, {
+		"within",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		func(root Node) Node {
+			return Search(root, `image "alpine"`)
+		},
+		2, 3,
+		true,
+	}, {
+		"within whitespace",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		func(root Node) Node {
+			return Search(root, `image "alpine"`)
+		},
+		2, 7,
+		true,
+	}, {
+		"not within",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		func(root Node) Node {
+			return Search(root, `image "alpine"`)
+		},
+		2, 1,
+		false,
+	}, {
+		"not within beyond EOL",
+		`
+		fs default() {
+			image "alpine"
+		}
+		`,
+		func(root Node) Node {
+			return Search(root, `image "alpine"`)
+		},
+		2, 99, // image "alpine" contains newline 2:2 -> 3:1
+		true,
+	}} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			mod := &Module{}
+			r := strings.NewReader(cleanup(tc.input))
+			err := Parser.Parse("", r, mod)
+			require.NoError(t, err)
+
+			n := tc.match(mod)
+			require.NotNil(t, n)
+
+			actual := IsPositionWithinNode(n, tc.line, tc.column)
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}