psi: optimize AST traversal and fix tree-sitter resource leaks (#178)

Author: wenxuanjun

src/analyzer/index/IndexingRoot.v

@@ -178,9 +178,9 @@ pub fn (mut i IndexingRoot) index() BuiltIndexStatus {
 	return .from_scratch
 }
 
-pub fn (mut i IndexingRoot) index_file(path string, content string) !FileIndex {
+pub fn (mut i IndexingRoot) index_file(path string, content string, mut p parser.Parser) !FileIndex {
 	last_modified := os.file_last_mod_unix(path)
-	res := parser.parse_code(content)
+	res := p.parse_code(content)
 	psi_file := psi.new_psi_file(path, res.tree, content)
 	module_fqn := psi.module_qualified_name(psi_file, i.root)
 
@@ -221,6 +221,8 @@ pub fn (mut i IndexingRoot) spawn_indexing_workers(cache_chan chan FileIndex, fi
 	wg.add(workers)
 	for j := 0; j < workers; j++ {
 		spawn fn [file_chan, mut wg, mut i, cache_chan] () {
+			mut p := parser.Parser.new()
+			defer { p.free() }
 			for {
 				filepath := <-file_chan or { break }
 				content := os.read_file(filepath) or {
@@ -230,7 +232,7 @@ pub fn (mut i IndexingRoot) spawn_indexing_workers(cache_chan chan FileIndex, fi
 					}).error('Error reading file for index')
 					continue
 				}
-				cache_chan <- i.index_file(filepath, content) or {
+				cache_chan <- i.index_file(filepath, content, mut p) or {
 					loglib.with_fields({
 						'uri':   lsp.document_uri_from_path(filepath).str()
 						'error': err.str()
@@ -303,7 +305,10 @@ pub fn (mut i IndexingRoot) mark_as_dirty(filepath string, new_content string) !
 		'uri': lsp.document_uri_from_path(filepath).str()
 	}).info('Marking document as dirty')
 	i.index.per_file.data.delete(filepath)
-	res := i.index_file(filepath, new_content) or {
+
+	mut p := parser.Parser.new()
+	defer { p.free() }
+	res := i.index_file(filepath, new_content, mut p) or {
 		return error('Error indexing dirty ${filepath}: ${err}')
 	}
 	i.index.per_file.data[filepath] = res
@@ -321,7 +326,9 @@ pub fn (mut i IndexingRoot) add_file(filepath string, content string) !FileIndex
 		'uri': lsp.document_uri_from_path(filepath).str()
 	}).info('Adding new document')
 
-	res := i.index_file(filepath, content) or {
+	mut p := parser.Parser.new()
+	defer { p.free() }
+	res := i.index_file(filepath, content, mut p) or {
 		return error('Error indexing added ${filepath}: ${err}')
 	}
 	i.index.per_file.data[filepath] = res

src/analyzer/index/StubTree.v

@@ -55,35 +55,56 @@ pub fn (tree &StubTree) get_imported_modules() []string {
 }
 
 pub fn build_stub_tree(file &psi.PsiFile, indexing_root string) &StubTree {
-	root := file.root()
+	mut walker := psi.new_tree_walker(file.tree.root_node())
+	defer { walker.free() }
+
 	stub_root := psi.new_root_stub(file.path())
 	module_fqn := psi.module_qualified_name(file, indexing_root)
 
-	build_stub_tree_for_node(root, stub_root, module_fqn, false)
+	if walker.current_node() != none {
+		build_stub_tree_recurse(mut walker, file, stub_root, module_fqn, false)
+	}
 
 	return &StubTree{
 		root: stub_root
 	}
 }
 
-pub fn build_stub_tree_for_node(node psi.PsiElement, parent &psi.StubBase, module_fqn string, build_for_all_children bool) {
-	element_type := psi.StubbedElementType{}
+fn build_stub_tree_recurse(mut tw psi.TreeWalker, file &psi.PsiFile, parent &psi.StubBase, module_fqn string, build_for_all_children bool) {
+	node := tw.current_node() or { return }
+	node_type := node.type_name
+
+	stub_type := psi.node_type_to_stub_type(node_type)
+	is_stubbable := stub_type != .root || psi.node_is_type(node_type)
 
-	node_copy := node
+	mut effective_parent := unsafe { parent }
+	mut should_traverse_children := true
+	mut pass_down_build_all := false
 
-	if node_copy is psi.StubBasedPsiElement || psi.node_is_type(node) || build_for_all_children {
-		if stub := element_type.create_stub(node, parent, module_fqn) {
-			is_qualified_type := node is psi.QualifiedType
-			for child in node.children() {
-				build_stub_tree_for_node(child, stub, module_fqn, build_for_all_children
-					|| is_qualified_type)
+	if is_stubbable || build_for_all_children {
+		psi_element := psi.create_element(node, file)
+		element_type := psi.StubbedElementType{}
+
+		if stub := element_type.create_stub(psi_element, parent, module_fqn) {
+			effective_parent = unsafe { stub }
+			if node_type == .qualified_type {
+				pass_down_build_all = true
 			}
 		}
-		return
 	}
 
-	for child in node.children() {
-		build_stub_tree_for_node(child, parent, module_fqn, false)
+	if should_traverse_children {
+		if tw.to_first_child() {
+			for {
+				build_stub_tree_recurse(mut tw, file, effective_parent, module_fqn,
+					build_for_all_children || pass_down_build_all)
+
+				if !tw.next_sibling() {
+					break
+				}
+			}
+			tw.to_parent()
+		}
 	}
 }
 

src/analyzer/parser/batch.v

@@ -32,15 +32,17 @@ pub fn parse_batch_files(files []string, count_workers int) []ParseResult {
 fn spawn_parser_workers(result_chan chan ParseResult, file_chan chan string, count_workers int) {
 	mut wg := sync.new_waitgroup()
 	wg.add(count_workers)
+
 	for i := 0; i < count_workers; i++ {
 		spawn fn [file_chan, mut wg, result_chan] () {
+			mut p := Parser.new()
+			defer { p.free() }
 			for {
 				filepath := <-file_chan or { break }
-				mut result := parse_file(filepath) or { continue }
+				mut result := p.parse_file(filepath) or { continue }
 				result.path = filepath
 				result_chan <- result
 			}
-
 			wg.done()
 		}()
 	}

src/analyzer/parser/parser.v

@@ -15,6 +15,28 @@ pub mut:
 // Source represent the possible types of V source code to parse.
 type Source = []byte | string
 
+// Parser is a wrapper around the Tree-sitter V parser.
+pub struct Parser {
+mut:
+	binding_parser &bindings.Parser[bindings.NodeType] = unsafe { nil }
+}
+
+// new creates a new Parser instance.
+pub fn Parser.new() &Parser {
+	mut bp := bindings.new_parser[bindings.NodeType](bindings.type_factory)
+	bp.set_language(bindings.language)
+	return &Parser{
+		binding_parser: bp
+	}
+}
+
+// free frees the Tree-sitter parser.
+pub fn (p &Parser) free() {
+	unsafe {
+		p.binding_parser.free()
+	}
+}
+
 // parse_file parses a V source file and returns the corresponding `tree_sitter.Tree` and `Rope`.
 // If the file could not be read, an error is returned.
 // If the file was read successfully, but could not be parsed, the result
@@ -25,18 +47,19 @@ type Source = []byte | string
 // import parser
 //
 // fn main() {
-//  res := parser.parse_file('foo.v') or {
+//  mut p := parser.Parser.new()
+//  res := p.parse_file('foo.v') or {
 //     eprintln('Error: could not parse file: ${err}')
 //     return
 //   }
 //   println(res.tree)
 // }
 // ```
-pub fn parse_file(filename string) !ParseResult {
-	mut file := os.read_file(filename) or {
-		return error('could not read file ${filename}: ${err}')
-	}
-	return parse_source(file)
+pub fn (mut p Parser) parse_file(filename string) !ParseResult {
+	content := os.read_file(filename) or { return error('could not read file ${filename}: ${err}') }
+	mut res := p.parse_source(content)
+	res.path = filename
+	return res
 }
 
 // parse_source parses a V code and returns the corresponding `tree_sitter.Tree` and `Rope`.
@@ -48,14 +71,15 @@ pub fn parse_file(filename string) !ParseResult {
 // import parser
 //
 // fn main() {
-//   res := parser.parse_source('fn main() { println("Hello, World!") }') or {
+//   mut p := parser.Parser.new()
+//   res := p.parse_source('fn main() { println("Hello, World!") }') or {
 //     eprintln('Error: could not parse source: ${err}')
 //     return
 //   }
 //   println(res.tree)
 // }
 // ```
-pub fn parse_source(source Source) ParseResult {
+pub fn (mut p Parser) parse_source(source Source) ParseResult {
 	code := match source {
 		string {
 			source
@@ -64,14 +88,18 @@ pub fn parse_source(source Source) ParseResult {
 			source.str()
 		}
 	}
-	return parse_code(code)
+	return p.parse_code(code)
 }
 
 // parse_code parses a V code and returns the corresponding `tree_sitter.Tree` and `Rope`.
 // Unlike `parse_file` and `parse_source`, `parse_code` don't return an error since
 // the source is always valid.
-pub fn parse_code(code string) ParseResult {
-	return parse_code_with_tree(code, unsafe { nil })
+pub fn (mut p Parser) parse_code(code string) ParseResult {
+	tree := p.binding_parser.parse_string(source: code)
+	return ParseResult{
+		tree:        tree
+		source_text: code
+	}
 }
 
 // parse_code_with_tree parses a V code and returns the corresponding `tree_sitter.Tree` and `Rope`.
@@ -86,19 +114,18 @@ pub fn parse_code(code string) ParseResult {
 // import parser
 //
 // fn main() {
+//   mut p := parser.Parser.new()
 //   code := 'fn main() { println("Hello, World!") }'
-//   res := parser.parse_code_with_tree(code, unsafe { nil })
+//   res := p.parse_code_with_tree(code, unsafe { nil })
 //   println(res.tree)
 //   // some changes in code
 //   code2 := 'fn foo() { println("Hello, World!") }'
-//   res2 = parser.parse_code_with_tree(code2, res.tree)
+//   res2 = p.parse_code_with_tree(code2, res.tree)
 //   println(res2.tree
 // }
-pub fn parse_code_with_tree(code string, old_tree &bindings.Tree[bindings.NodeType]) ParseResult {
-	mut parser := bindings.new_parser[bindings.NodeType](bindings.type_factory)
-	parser.set_language(bindings.language)
+pub fn (mut p Parser) parse_code_with_tree(code string, old_tree &bindings.Tree[bindings.NodeType]) ParseResult {
 	raw_tree := if isnil(old_tree) { unsafe { nil } } else { old_tree.raw_tree }
-	tree := parser.parse_string(source: code, tree: raw_tree)
+	tree := p.binding_parser.parse_string(source: code, tree: raw_tree)
 	return ParseResult{
 		tree:        tree
 		source_text: code

src/analyzer/psi/ConstantDefinition.v

@@ -14,7 +14,13 @@ pub fn (c &ConstantDefinition) is_public() bool {
 
 pub fn (c &ConstantDefinition) get_type() types.Type {
 	expr := c.expression() or { return types.unknown_type }
-	return infer_type(expr)
+	res := infer_type(expr)
+	if c.stub_based() {
+		if mut file := expr.containing_file() {
+			file.free()
+		}
+	}
+	return res
 }
 
 fn (c &ConstantDefinition) identifier() ?PsiElement {
@@ -68,7 +74,9 @@ pub fn (c &ConstantDefinition) expression() ?PsiElement {
 	if stub := c.get_stub() {
 		file := c.containing_file() or { return none }
 		// pretty hacky but it works
-		res := parser.parse_code(stub.additional)
+		mut p := parser.Parser.new()
+		defer { p.free() }
+		res := p.parse_code(stub.additional)
 		root := res.tree.root_node()
 		first_child := root.first_child()?
 		next_first_child := first_child.first_child()?

src/analyzer/psi/EnumFieldDeclaration.v

@@ -66,7 +66,9 @@ pub fn (f &EnumFieldDeclaration) value() ?PsiElement {
 		}
 
 		file := f.containing_file() or { return none }
-		res := parser.parse_code(stub.additional)
+		mut p := parser.Parser.new()
+		defer { p.free() }
+		res := p.parse_code(stub.additional)
 		root := res.tree.root_node()
 		first_child := root.first_child()?
 		next_first_child := first_child.first_child()?
@@ -133,7 +135,13 @@ fn (f &EnumFieldDeclaration) get_value_impl() i64 {
 
 	if !is_flag {
 		if value := f.value() {
-			if val := f.calculate_value(value) {
+			val := f.calculate_value(value)
+			if f.stub_based() {
+				if mut file := value.containing_file() {
+					file.free()
+				}
+			}
+			if val != none {
 				return val
 			}
 		}