v.ast: add walker submodule for ast walking (#7775)

vlib/v/ast/ast.v

@@ -25,6 +25,11 @@ pub type Stmt = AssertStmt | AssignStmt | Block | BranchStmt | CompFor | ConstDe
 // the .position() token.Position methods too.
 pub type ScopeObject = ConstField | GlobalField | Var
 
+// TOOD: replace table.Param
+pub type Node = ConstField | EnumField | Expr | Field | File | GlobalField | IfBranch |
+	MatchBranch | ScopeObject | SelectBranch | Stmt | StructField | StructInitField | table.Param
+	
+
 pub struct Type {
 pub:
 	typ table.Type
@@ -134,6 +139,7 @@ pub:
 	name       string
 	attrs      []table.Attr
 	pos        token.Position
+	name_pos   token.Position // `name` in import name
 	is_skipped bool // module main can be skipped in single file programs
 }
 
@@ -269,6 +275,8 @@ pub:
 	mod       string // the module name of the import
 	alias     string // the `x` in `import xxx as x`
 	pos       token.Position
+	mod_pos   token.Position
+	alias_pos token.Position
 pub mut:
 	syms      []ImportSymbol // the list of symbols in `import {symbol1, symbol2}`
 }
@@ -1257,6 +1265,180 @@ pub fn (stmt Stmt) position() token.Position {
 	}
 }
 
+pub fn (node Node) position() token.Position {
+	match node {
+		Stmt {
+			mut pos := node.position()
+			if node is Import {
+				for sym in node.syms {
+					pos = pos.extend(sym.pos)
+				}
+			}
+			return pos
+		}
+		Expr {
+			return node.position()
+		}
+		StructField {
+			return node.pos.extend(node.type_pos)
+		}
+		MatchBranch, SelectBranch, Field, EnumField, ConstField, StructInitField, GlobalField, table.Param {
+			return node.pos
+		}
+		IfBranch {
+			return node.pos.extend(node.body_pos)
+		}
+		ScopeObject {
+			match node {
+				ConstField, GlobalField, Var { return node.pos }
+			}
+		}
+		File {
+			mut pos := token.Position{}
+			if node.stmts.len > 0 {
+				first_pos := node.stmts.first().position()
+				last_pos := node.stmts.last().position()
+				pos = first_pos.extend_with_last_line(last_pos, last_pos.line_nr)
+			}
+			return pos
+		}
+	}
+}
+
+pub fn (node Node) children() []Node {
+	mut children := []Node{}
+	if node is Expr {
+		match node {
+			StringInterLiteral, Assoc, ArrayInit {
+				return node.exprs.map(Node(it))
+			}
+			SelectorExpr, PostfixExpr, UnsafeExpr, AsCast, ParExpr, IfGuardExpr, SizeOf, Likely, TypeOf, ArrayDecompose {
+				children << node.expr
+			}
+			LockExpr, OrExpr {
+				return node.stmts.map(Node(it))
+			}
+			StructInit {
+				return node.fields.map(Node(it))
+			}
+			AnonFn {
+				children << Stmt(node.decl)
+			}
+			CallExpr {
+				children << node.left
+				children << Expr(node.or_block)
+			}
+			InfixExpr {
+				children << node.left
+				children << node.right
+			}
+			PrefixExpr {
+				children << node.right
+			}
+			IndexExpr {
+				children << node.left
+				children << node.index
+			}
+			IfExpr {
+				children << node.left
+				children << node.branches.map(Node(it))
+			}
+			MatchExpr {
+				children << node.cond
+				children << node.branches.map(Node(it))
+			}
+			SelectExpr {
+				return node.branches.map(Node(it))
+			}
+			ChanInit {
+				children << node.cap_expr
+			}
+			MapInit {
+				children << node.keys.map(Node(it))
+				children << node.vals.map(Node(it))
+			}
+			RangeExpr {
+				children << node.low
+				children << node.high
+			}
+			CastExpr {
+				children << node.expr
+				children << node.arg
+			}
+			ConcatExpr {
+				return node.vals.map(Node(it))
+			}
+			ComptimeCall, ComptimeSelector {
+				children << node.left
+			}
+			else {}
+		}
+	} else if node is Stmt {
+		match node {
+			Block, DeferStmt, ForCStmt, ForInStmt, ForStmt, CompFor {
+				return node.stmts.map(Node(it))
+			}
+			ExprStmt, AssertStmt {
+				children << node.expr
+			}
+			InterfaceDecl {
+				return node.methods.map(Node(Stmt(it)))
+			}
+			AssignStmt {
+				children << node.left.map(Node(it))
+				children << node.right.map(Node(it))
+			}
+			Return {
+				return node.exprs.map(Node(it))
+			}
+			// NB: these four decl nodes cannot be merged as one branch
+			StructDecl {
+				return node.fields.map(Node(it))
+			}
+			GlobalDecl {
+				return node.fields.map(Node(it))
+			}
+			ConstDecl {
+				return node.fields.map(Node(it))
+			}
+			EnumDecl {
+				return node.fields.map(Node(it))
+			}
+			FnDecl {
+				if node.is_method {
+					children << Node(node.receiver)
+				}
+				children << node.params.map(Node(it))
+				children << node.stmts.map(Node(it))
+			}
+			else {}
+		}
+	} else if node is ScopeObject {
+		match node {
+			GlobalField, ConstField, Var { children << node.expr }
+		}
+	} else {
+		match node {
+			GlobalField, ConstField, EnumField, StructInitField {
+				children << node.expr
+			}
+			SelectBranch {
+				children << node.stmt
+				children << node.stmts.map(Node(it))
+			}
+			IfBranch, File {
+				return node.stmts.map(Node(it))
+			}
+			MatchBranch {
+				children << node.stmts.map(Node(it))
+				children << node.exprs.map(Node(it))
+			}
+			else {}
+		}
+	}
+	return children
+}
+
 // TODO: remove this fugly hack :-|
 // fe2ex/1 and ex2fe/1 are used to convert back and forth from
 // table.FExpr to ast.Expr , which in turn is needed to break

vlib/v/ast/walker/walker.v

@@ -0,0 +1,37 @@
+module walker
+
+import v.ast
+
+// Visitor defines a visit method which is invoked by the walker in each node it encounters.
+pub interface Visitor {
+	visit(node ast.Node) ?
+}
+
+pub type InspectorFn = fn (node ast.Node, data voidptr) bool
+
+struct Inspector {
+	inspector_callback InspectorFn
+mut:
+	data               voidptr
+}
+
+pub fn (i &Inspector) visit(node ast.Node) ? {
+	if i.inspector_callback(node, i.data) {
+		return
+	}
+	return none
+}
+
+// inspect traverses and checks the AST node on a depth-first order and based on the data given
+pub fn inspect(node ast.Node, data voidptr, inspector_callback InspectorFn) {
+	walk(Inspector{inspector_callback, data}, node)
+}
+
+// walk traverses the AST using the given visitor
+pub fn walk(visitor Visitor, node ast.Node) {
+	visitor.visit(node) or { return }
+	children := node.children()
+	for child_node in children {
+		walk(visitor, &child_node)
+	}
+}

vlib/v/ast/walker/walker_test.v

@@ -0,0 +1,71 @@
+import v.ast
+import v.ast.walker
+import v.parser
+import v.table
+import v.pref
+
+fn parse_text(text string) ast.File {
+	tbl := table.new_table()
+	prefs := pref.new_preferences()
+	scope := &ast.Scope{
+		parent: 0
+	}
+	return parser.parse_text(text, '', tbl, .skip_comments, prefs, scope)
+}
+
+struct NodeByOffset {
+	pos  int
+mut:
+	node ast.Node
+}
+
+fn (mut n NodeByOffset) visit(node ast.Node) ? {
+	node_pos := node.position()
+	if n.pos >= node_pos.pos && n.pos <= node_pos.pos + node_pos.len && node !is ast.File {
+		n.node = node
+		return none
+	}
+	return
+}
+
+fn test_walk() {
+	source := '
+module main
+struct Foo {
+	name string
+}
+	'
+	file := parse_text(source)
+	mut nbo := NodeByOffset{
+		pos: 13
+	}
+	walker.walk(nbo, file)
+	assert nbo.node is ast.Stmt
+	stmt := nbo.node as ast.Stmt
+	assert stmt is ast.StructDecl
+}
+
+fn test_inspect() {
+	source := '
+module main
+	'
+	file := parse_text(source)
+	walker.inspect(&file, voidptr(0), fn (node ast.Node, data voidptr) bool {
+		// Second visit must be ast.Stmt
+		if node is ast.Stmt {
+			if node !is ast.Module {
+				// Proceed to another node
+				return false
+			}
+			assert node is ast.Module
+			mod := node as ast.Module
+			assert mod.name == 'main'
+			return false
+		}
+		// First visit must be ast.File
+		assert node is ast.File
+		// True means that the inspector must now
+		// inspect the ast.File's children
+		return true
+	})
+}

vlib/v/checker/checker.v

@@ -90,7 +90,7 @@ pub fn (mut c Checker) check(ast_file &ast.File) {
 	for i, ast_import in ast_file.imports {
 		for j in 0 .. i {
 			if ast_import.mod == ast_file.imports[j].mod {
-				c.error('module name `$ast_import.mod` duplicate', ast_import.pos)
+				c.error('module name `$ast_import.mod` duplicate', ast_import.mod_pos)
 			}
 		}
 	}

vlib/v/parser/module.v

@@ -60,7 +60,7 @@ fn (mut p Parser) check_unused_imports() {
 		mod := import_m.mod
 		if !p.is_used_import(alias) {
 			mod_alias := if alias == mod { alias } else { '$alias ($mod)' }
-			p.warn_with_pos("module '$mod_alias' is imported but never used", import_m.pos)
+			p.warn_with_pos("module '$mod_alias' is imported but never used", import_m.mod_pos)
 		}
 	}
 }

vlib/v/parser/parser.v

@@ -1683,27 +1683,35 @@ fn (mut p Parser) module_decl() ast.Module {
 	mut name := 'main'
 	is_skipped := p.tok.kind != .key_module
 	mut module_pos := token.Position{}
+	mut name_pos := token.Position{}
+	mut mod_node := ast.Module{}
 	if !is_skipped {
 		p.attrs = []
 		module_pos = p.tok.position()
 		p.next()
-		mut pos := p.tok.position()
+		name_pos = p.tok.position()
 		name = p.check_name()
-		if module_pos.line_nr != pos.line_nr {
-			p.error_with_pos('`module` and `$name` must be at same line', pos)
-			return ast.Module{}
+		mod_node = ast.Module{
+			pos: module_pos
 		}
-		pos = p.tok.position()
-		if module_pos.line_nr == pos.line_nr && p.tok.kind != .comment {
+		if module_pos.line_nr != name_pos.line_nr {
+			p.error_with_pos('`module` and `$name` must be at same line', name_pos)
+			return mod_node
+		}
+		// NB: this shouldn't be reassigned into name_pos
+		// as it creates a wrong position when extended
+		// to module_pos
+		n_pos := p.tok.position()
+		if module_pos.line_nr == n_pos.line_nr && p.tok.kind != .comment {
 			if p.tok.kind != .name {
-				p.error_with_pos('`module x` syntax error', pos)
-				return ast.Module{}
+				p.error_with_pos('`module x` syntax error', n_pos)
+				return mod_node
 			} else {
-				p.error_with_pos('`module x` can only declare one module', pos)
-				return ast.Module{}
+				p.error_with_pos('`module x` can only declare one module', n_pos)
+				return mod_node
 			}
 		}
-		module_pos = module_pos.extend(pos)
+		module_pos = module_pos.extend(name_pos)
 	}
 	mut full_mod := p.table.qualify_module(name, p.file_name)
 	if p.pref.build_mode == .build_module && !full_mod.contains('.') {
@@ -1721,6 +1729,13 @@ fn (mut p Parser) module_decl() ast.Module {
 	}
 	p.mod = full_mod
 	p.builtin_mod = p.mod == 'builtin'
+	mod_node = ast.Module{
+		name: full_mod
+		attrs: module_attrs
+		is_skipped: is_skipped
+		pos: module_pos
+		name_pos: name_pos
+	}
 	if !is_skipped {
 		for ma in module_attrs {
 			match ma.name {
@@ -1729,78 +1744,103 @@ fn (mut p Parser) module_decl() ast.Module {
 				}
 				else {
 					p.error_with_pos('unknown module attribute `[$ma.name]`', ma.pos)
-					return ast.Module{}
+					return mod_node
 				}
 			}
 		}
 	}
-	return ast.Module{
-		name: full_mod
-		attrs: module_attrs
-		is_skipped: is_skipped
-		pos: module_pos
-	}
+	return mod_node
 }
 
 fn (mut p Parser) import_stmt() ast.Import {
 	import_pos := p.tok.position()
 	p.check(.key_import)
-	pos := p.tok.position()
+	mut pos := p.tok.position()
+	mut import_node := ast.Import{
+		pos: import_pos.extend(pos)
+	}
 	if p.tok.kind == .lpar {
 		p.error_with_pos('`import()` has been deprecated, use `import x` instead', pos)
-		return ast.Import{}
+		return import_node
 	}
-	mut mod_name := p.check_name()
+	mut mod_name_arr := []string{}
+	mod_name_arr << p.check_name()
 	if import_pos.line_nr != pos.line_nr {
 		p.error_with_pos('`import` statements must be a single line', pos)
-		return ast.Import{}
+		return import_node
+	}
+	mut mod_alias := mod_name_arr[0]
+	import_node = ast.Import{
+		pos: import_pos.extend(pos)
+		mod_pos: pos
+		alias_pos: pos
 	}
-	mut mod_alias := mod_name
 	for p.tok.kind == .dot {
 		p.next()
-		pos_t := p.tok.position()
+		submod_pos := p.tok.position()
 		if p.tok.kind != .name {
-			p.error_with_pos('module syntax error, please use `x.y.z`', pos)
-			return ast.Import{}
+			p.error_with_pos('module syntax error, please use `x.y.z`', submod_pos)
+			return import_node
 		}
-		if import_pos.line_nr != pos_t.line_nr {
-			p.error_with_pos('`import` and `submodule` must be at same line', pos)
-			return ast.Import{}
+		if import_pos.line_nr != submod_pos.line_nr {
+			p.error_with_pos('`import` and `submodule` must be at same line', submod_pos)
+			return import_node
 		}
 		submod_name := p.check_name()
-		mod_name += '.' + submod_name
+		mod_name_arr << submod_name
 		mod_alias = submod_name
-	}
-	if p.tok.kind == .key_as {
-		p.next()
-		mod_alias = p.check_name()
-		if mod_alias == mod_name.split('.').last() {
-			p.error_with_pos('import alias `$mod_name as $mod_alias` is redundant', p.prev_tok.position())
-			return ast.Import{}
-		}
-	}
-	mut node := ast.Import{
-		pos: pos
-		mod: mod_name
+		pos = pos.extend(submod_pos)
+		import_node = ast.Import{
+			pos: import_pos.extend(pos)
+			mod_pos: pos
+			alias_pos: submod_pos
+			mod: mod_name_arr.join('.')
 			alias: mod_alias
 		}
+	}
+	if mod_name_arr.len == 1 {
+		import_node = ast.Import{
+			pos: import_node.pos
+			mod_pos: import_node.mod_pos
+			alias_pos: import_node.alias_pos
+			mod: mod_name_arr[0]
+			alias: mod_alias
+		}
+	}
+	mod_name := import_node.mod
+	if p.tok.kind == .key_as {
+		p.next()
+		alias_pos := p.tok.position()
+		mod_alias = p.check_name()
+		if mod_alias == mod_name_arr.last() {
+			p.error_with_pos('import alias `$mod_name as $mod_alias` is redundant', p.prev_tok.position())
+			return import_node
+		}
+		import_node = ast.Import{
+			pos: import_node.pos.extend(alias_pos)
+			mod_pos: import_node.mod_pos
+			alias_pos: alias_pos
+			mod: import_node.mod
+			alias: mod_alias
+		}
+	}
 	if p.tok.kind == .lcbr { // import module { fn1, Type2 } syntax
-		p.import_syms(mut node)
+		p.import_syms(mut import_node)
 		p.register_used_import(mod_alias) // no `unused import` msg for parent
 	}
 	pos_t := p.tok.position()
 	if import_pos.line_nr == pos_t.line_nr {
 		if p.tok.kind !in [.lcbr, .eof, .comment] {
 			p.error_with_pos('cannot import multiple modules at a time', pos_t)
-			return ast.Import{}
+			return import_node
 		}
 	}
 	p.imports[mod_alias] = mod_name
 	// if mod_name !in p.table.imports {
 	p.table.imports << mod_name
-	p.ast_imports << node
+	p.ast_imports << import_node
 	// }
-	return node
+	return import_node
 }
 
 // import_syms parses the inner part of `import module { submod1, submod2 }`