optimize build symbol table

cmd/tlin/main.go

@@ -43,6 +43,7 @@ func main() {
 		os.Exit(1)
 	}
 
+	// TODO: Cache the directory tree to avoid re-traversing the same directories.
 	rootDir := "."
 	engine, err := internal.NewEngine(rootDir)
 	if err != nil {
@@ -60,6 +61,7 @@ func main() {
 		}
 	}
 
+	// TODO: We might be use the cached directory result when execute analysis functions.
 	if *cyclomaticComplexity {
 		runWithTimeout(ctx, func() {
 			runCyclomaticComplexityAnalysis(args, *cyclomaticThreshold)

internal/symbol_table.go

@@ -4,47 +4,91 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
-	"os"
+	"io/fs"
+	_ "os"
 	"path/filepath"
+	"runtime"
 	"strings"
+	"sync"
 )
 
 type SymbolTable struct {
 	symbols map[string]string // symbol name -> file path
+	mu      sync.RWMutex
 }
 
 func BuildSymbolTable(rootDir string) (*SymbolTable, error) {
 	st := &SymbolTable{symbols: make(map[string]string)}
-	err := filepath.Walk(rootDir, func(path string, info os.FileInfo, err error) error {
-		if err != nil {
-			return err
-		}
-		if !info.IsDir() && (strings.HasSuffix(path, ".go") || strings.HasSuffix(path, ".gno")) {
-			if err := st.parseFile(path); err != nil {
+	errChan := make(chan error, 1)
+	fileChan := make(chan string, 100)
+
+	var wg sync.WaitGroup
+	for i := 0; i < runtime.NumCPU(); i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for path := range fileChan {
+				if err := st.parseFile(path); err != nil {
+					select {
+					case errChan <- err:
+					default:
+					}
+					return
+				}
+			}
+		}()
+	}
+
+	go func() {
+		err := filepath.WalkDir(rootDir, func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
 				return err
 			}
+			if !d.IsDir() && (strings.HasSuffix(path, ".go") || strings.HasSuffix(path, ".gno")) {
+				fileChan <- path
+			}
+			return nil
+		})
+		close(fileChan)
+		if err != nil {
+			errChan <- err
 		}
-		return err
-	})
-	return st, err
+	}()
+
+	wg.Wait()
+
+	select {
+	case err := <-errChan:
+		return nil, err
+	default:
+		return st, nil
+	}
 }
 
-func (st *SymbolTable) parseFile(filepath string) error {
+func (st *SymbolTable) parseFile(path string) error {
 	fset := token.NewFileSet()
-	node, err := parser.ParseFile(fset, filepath, nil, parser.AllErrors)
+	node, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
 	if err != nil {
 		return err
 	}
 
+	packageName := node.Name.Name
 	ast.Inspect(node, func(n ast.Node) bool {
 		switch x := n.(type) {
-		case *ast.TypeSpec:
-			st.symbols[x.Name.Name] = filepath
 		case *ast.FuncDecl:
-			st.symbols[x.Name.Name] = filepath
+			funcName := x.Name.Name
+			fullName := packageName + "." + funcName
+			st.addSymbol(fullName, path)
+
+		case *ast.TypeSpec:
+			typeName := x.Name.Name
+			fullName := packageName + "." + typeName
+			st.addSymbol(fullName, path)
 		case *ast.ValueSpec:
 			for _, ident := range x.Names {
-				st.symbols[ident.Name] = filepath
+				varName := ident.Name
+				fullName := packageName + "." + varName
+				st.addSymbol(fullName, path)
 			}
 		}
 		return true
@@ -54,11 +98,24 @@ func (st *SymbolTable) parseFile(filepath string) error {
 }
 
 func (st *SymbolTable) IsDefined(symbol string) bool {
+	st.mu.RLock()
+	defer st.mu.RUnlock()
+
 	_, exists := st.symbols[symbol]
 	return exists
 }
 
 func (st *SymbolTable) GetSymbolPath(symbol string) (string, bool) {
+	st.mu.RLock()
+	defer st.mu.RUnlock()
+
 	path, exists := st.symbols[symbol]
 	return path, exists
 }
+
+func (st *SymbolTable) addSymbol(key, value string) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+
+	st.symbols[key] = value
+}

internal/symbol_table_test.go

@@ -1,8 +1,10 @@
 package internal
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -20,36 +22,119 @@ type TestStruct struct {}
 func TestFunc() {}
 var TestVar int
 `
-	err = os.WriteFile(filepath.Join(tmpDir, "file1.go"), []byte(file1Content), 0o644)
+	file1Path := filepath.Join(tmpDir, "file1.go")
+	err = os.WriteFile(file1Path, []byte(file1Content), 0o644)
 	require.NoError(t, err)
 
 	file2Content := `package test
 type AnotherStruct struct {}
 func AnotherFunc() {}
 `
-	err = os.WriteFile(filepath.Join(tmpDir, "file2.go"), []byte(file2Content), 0o644)
+	file2Path := filepath.Join(tmpDir, "file2.go")
+	err = os.WriteFile(file2Path, []byte(file2Content), 0o644)
 	require.NoError(t, err)
 
-	// create symbol table
 	st, err := BuildSymbolTable(tmpDir)
 	require.NoError(t, err)
 
-	assert.True(t, st.IsDefined("TestStruct"))
-	assert.True(t, st.IsDefined("TestFunc"))
-	assert.True(t, st.IsDefined("TestVar"))
-	assert.True(t, st.IsDefined("AnotherStruct"))
-	assert.True(t, st.IsDefined("AnotherFunc"))
-	assert.False(t, st.IsDefined("NonExistentSymbol"))
+	assert.True(t, st.IsDefined("test.TestStruct"))
+	assert.True(t, st.IsDefined("test.TestFunc"))
+	assert.True(t, st.IsDefined("test.TestVar"))
+	assert.True(t, st.IsDefined("test.AnotherStruct"))
+	assert.True(t, st.IsDefined("test.AnotherFunc"))
+	assert.False(t, st.IsDefined("test.NonExistentSymbol"))
 
 	// validate symbol file paths
-	path, exists := st.GetSymbolPath("TestStruct")
+	path, exists := st.GetSymbolPath("test.TestStruct")
 	assert.True(t, exists)
-	assert.Equal(t, filepath.Join(tmpDir, "file1.go"), path)
+	assert.Equal(t, file1Path, path)
 
-	path, exists = st.GetSymbolPath("AnotherFunc")
+	path, exists = st.GetSymbolPath("test.AnotherFunc")
 	assert.True(t, exists)
-	assert.Equal(t, filepath.Join(tmpDir, "file2.go"), path)
+	assert.Equal(t, file2Path, path)
 
-	_, exists = st.GetSymbolPath("NonExistentSymbol")
+	_, exists = st.GetSymbolPath("test.NonExistentSymbol")
 	assert.False(t, exists)
 }
+
+func TestConcurrentSymbolTable(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "concurrent-symboltable-test")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	for i := 0; i < 100; i++ {
+		content := fmt.Sprintf(`package test%d
+func Func%d() {}
+var Var%d int
+`, i, i, i)
+		err = os.WriteFile(filepath.Join(tmpDir, fmt.Sprintf("file%d.go", i)), []byte(content), 0o644)
+		require.NoError(t, err)
+	}
+
+	st, err := BuildSymbolTable(tmpDir)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			assert.True(t, st.IsDefined(fmt.Sprintf("test%d.Func%d", i, i)))
+			assert.True(t, st.IsDefined(fmt.Sprintf("test%d.Var%d", i, i)))
+		}(i)
+	}
+	wg.Wait()
+}
+
+func BenchmarkBuildSymbolTable(b *testing.B) {
+	tmpDir, err := os.MkdirTemp("", "benchmark-symboltable")
+	require.NoError(b, err)
+	defer os.RemoveAll(tmpDir)
+
+	numFiles := 1000
+	for i := 0; i < numFiles; i++ {
+		content := fmt.Sprintf(`package test%d
+func Func%d() {}
+var Var%d int
+type Struct%d struct{}
+`, i, i, i, i)
+		err = os.WriteFile(filepath.Join(tmpDir, fmt.Sprintf("file%d.go", i)), []byte(content), 0o644)
+		require.NoError(b, err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := BuildSymbolTable(tmpDir)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkSymbolTableIsDefined(b *testing.B) {
+	tmpDir, err := os.MkdirTemp("", "benchmark-symboltable-isdefined")
+	require.NoError(b, err)
+	defer os.RemoveAll(tmpDir)
+
+	numFiles := 1000
+	for i := 0; i < numFiles; i++ {
+		content := fmt.Sprintf(`package test%d
+func Func%d() {}
+var Var%d int
+type Struct%d struct{}
+`, i, i, i, i)
+		err = os.WriteFile(filepath.Join(tmpDir, fmt.Sprintf("file%d.go", i)), []byte(content), 0o644)
+		require.NoError(b, err)
+	}
+
+	st, err := BuildSymbolTable(tmpDir)
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		randomIndex := i % numFiles
+		st.IsDefined(fmt.Sprintf("test%d.Func%d", randomIndex, randomIndex))
+		st.IsDefined(fmt.Sprintf("test%d.Var%d", randomIndex, randomIndex))
+		st.IsDefined(fmt.Sprintf("test%d.Struct%d", randomIndex, randomIndex))
+	}
+}