feat: speed up build graph (cloudwego#156)

Author: wzekin

lang/uniast/ast_test.go

@@ -42,3 +42,17 @@ func TestRepository_BuildGraph(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkRepository_BuildGraph(b *testing.B) {
+	astFile := testutils.GetTestAstFile("large_ast")
+	r, err := LoadRepo(astFile)
+	if err != nil {
+		b.Fatalf("failed to load repo: %v", err)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := r.BuildGraph(); err != nil {
+			b.Fatalf("failed to build graph: %v", err)
+		}
+	}
+}

lang/uniast/node.go

@@ -159,14 +159,32 @@ func (r *Repository) AllNodesSetRepo() {
 }
 
 func (r *Repository) BuildGraph() error {
-	r.Graph = make(map[string]*Node)
+	// Optimization: Pre-calculate total number of internal nodes to pre-allocate Graph map
+	var totalNodes int
+	for _, mod := range r.Modules {
+		if mod.IsExternal() {
+			continue
+		}
+		for _, pkg := range mod.Packages {
+			totalNodes += len(pkg.Functions) + len(pkg.Types) + len(pkg.Vars)
+		}
+	}
+	r.Graph = make(map[string]*Node, totalNodes)
 	for _, mod := range r.Modules {
 		if mod.IsExternal() {
 			continue
 		}
 		for _, pkg := range mod.Packages {
 			for _, f := range pkg.Functions {
 				n := r.SetNode(f.Identity, FUNC)
+				// Pre-allocate Dependencies
+				capDeps := len(f.Params) + len(f.Results) + len(f.FunctionCalls) + len(f.MethodCalls) + len(f.Types) + len(f.GlobalVars)
+				if f.Receiver != nil {
+					capDeps++
+				}
+				if n.Dependencies == nil && capDeps > 0 {
+					n.Dependencies = make([]Relation, 0, capDeps)
+				}
 				for _, dep := range f.Params {
 					r.AddRelation(n, dep.Identity, dep.FileLine, DEPENDENCY)
 				}
@@ -193,6 +211,16 @@ func (r *Repository) BuildGraph() error {
 
 			for _, t := range pkg.Types {
 				n := r.SetNode(t.Identity, TYPE)
+				// Pre-allocate Dependencies, Inherits, Implements
+				if n.Dependencies == nil && len(t.SubStruct) > 0 {
+					n.Dependencies = make([]Relation, 0, len(t.SubStruct))
+				}
+				if n.Inherits == nil && len(t.InlineStruct) > 0 {
+					n.Inherits = make([]Relation, 0, len(t.InlineStruct))
+				}
+				if n.Implements == nil && len(t.Implements) > 0 {
+					n.Implements = make([]Relation, 0, len(t.Implements))
+				}
 				for _, dep := range t.SubStruct {
 					r.AddRelation(n, dep.Identity, dep.FileLine, DEPENDENCY)
 				}
@@ -206,6 +234,17 @@ func (r *Repository) BuildGraph() error {
 
 			for _, v := range pkg.Vars {
 				n := r.SetNode(v.Identity, VAR)
+				// Pre-allocate Dependencies and Groups
+				capDeps := len(v.Dependencies)
+				if v.Type != nil {
+					capDeps++
+				}
+				if n.Dependencies == nil && capDeps > 0 {
+					n.Dependencies = make([]Relation, 0, capDeps)
+				}
+				if n.Groups == nil && len(v.Groups) > 0 {
+					n.Groups = make([]Relation, 0, len(v.Groups))
+				}
 				if v.Type != nil {
 					r.AddRelation(n, *v.Type, v.FileLine, DEPENDENCY)
 				}

lang/uniast/utils.go

@@ -22,35 +22,35 @@ import (
 )
 
 func Append[T comparable](ids []T, id T) []T {
-	for _, i := range ids {
-		if i == id {
+	for i := range ids {
+		if ids[i] == id {
 			return ids
 		}
 	}
 	return append(ids, id)
 }
 
 func InsertDependency(ids []Dependency, id Dependency) []Dependency {
-	for _, i := range ids {
-		if i.Identity == id.Identity {
+	for i := range ids {
+		if ids[i].Identity == id.Identity {
 			return ids
 		}
 	}
 	return append(ids, id)
 }
 
 func InserImport(ids []Import, id Import) []Import {
-	for _, i := range ids {
-		if i.Path == id.Path {
+	for i := range ids {
+		if ids[i].Path == id.Path {
 			return ids
 		}
 	}
 	return append(ids, id)
 }
 
 func InsertRelation(ids []Relation, id Relation) []Relation {
-	for _, i := range ids {
-		if i.Identity == id.Identity {
+	for i := range ids {
+		if ids[i].Identity == id.Identity {
 			return ids
 		}
 	}