Merge pull request #5 from chinmay-sawant/fix-4-astref

Refactored the code as per the #4 changes

Author: chinmay-sawant

cmd/analyzer/utils.go

@@ -55,6 +55,9 @@ func FindFunctionBody(lines []string, funcLine int) (int, int) {
 func FindCalls(bodyLines []string) []string {
 	var calls []string
 	reCalls := regexp.MustCompile(`(\w+(?:\.\w+)*)\(`)
+	// Enhanced regex to capture function references passed as arguments
+	// This captures function names that appear as arguments (not followed by parentheses)
+	reFuncRefs := regexp.MustCompile(`(?:\s|,|\()(handle[A-Za-z]\w+)(?:\s*[,\)\s]|$)`)
 
 	// Standard library packages to ignore
 	standardPackages := map[string]bool{
@@ -140,7 +143,46 @@ func FindCalls(bodyLines []string) []string {
 		"Wait": true,
 	}
 
+	// Keywords and common variable names to exclude
+	excludeKeywords := map[string]bool{
+		"true":      true,
+		"false":     true,
+		"nil":       true,
+		"err":       true,
+		"error":     true,
+		"string":    true,
+		"int":       true,
+		"float":     true,
+		"bool":      true,
+		"byte":      true,
+		"rune":      true,
+		"if":        true,
+		"for":       true,
+		"switch":    true,
+		"case":      true,
+		"default":   true,
+		"return":    true,
+		"break":     true,
+		"continue":  true,
+		"goto":      true,
+		"var":       true,
+		"const":     true,
+		"type":      true,
+		"func":      true,
+		"package":   true,
+		"import":    true,
+		"range":     true,
+		"select":    true,
+		"go":        true,
+		"defer":     true,
+		"chan":      true,
+		"map":       true,
+		"struct":    true,
+		"interface": true,
+	}
+
 	for _, line := range bodyLines {
+		// Find traditional function calls (package.function())
 		matches := reCalls.FindAllStringSubmatch(line, -1)
 		for _, match := range matches {
 			call := match[1]
@@ -174,6 +216,34 @@ func FindCalls(bodyLines []string) []string {
 				calls = append(calls, call)
 			}
 		}
+
+		// Find function references passed as arguments
+		refMatches := reFuncRefs.FindAllStringSubmatch(line, -1)
+		for _, match := range refMatches {
+			if len(match) > 1 {
+				funcRef := strings.TrimSpace(match[1])
+
+				// Skip if it's a keyword or common variable name
+				if excludeKeywords[funcRef] {
+					continue
+				}
+
+				// Skip if it's too short to be meaningful
+				if len(funcRef) < 3 {
+					continue
+				}
+
+				// Skip if it contains dots (already handled by reCalls)
+				if strings.Contains(funcRef, ".") {
+					continue
+				}
+
+				// Add the current package prefix to make it consistent with other calls
+				if !contains(calls, funcRef) {
+					calls = append(calls, funcRef)
+				}
+			}
+		}
 	}
 	return calls
 }

cmd/main.go

@@ -93,6 +93,14 @@ func findFunctions(filePath, absPath, module string) ([]analyzer.FunctionInfo, e
 		}
 	}
 
+	// Collect all function names in this file for reference resolution
+	var localFunctions []string
+	for _, line := range lines {
+		if matches := re.FindStringSubmatch(line); matches != nil {
+			localFunctions = append(localFunctions, matches[1])
+		}
+	}
+
 	for i, line := range lines {
 		if matches := re.FindStringSubmatch(line); matches != nil {
 			funcInfo := analyzer.FunctionInfo{
@@ -104,7 +112,23 @@ func findFunctions(filePath, absPath, module string) ([]analyzer.FunctionInfo, e
 			start, end := analyzer.FindFunctionBody(lines, i)
 			if start != -1 && end != -1 && start+1 < end && end < len(lines) {
 				calls := analyzer.FindCalls(lines[start+1 : end])
-				funcInfo.Calls = calls
+
+				// Resolve local function references by adding package prefix
+				var resolvedCalls []string
+				for _, call := range calls {
+					if !strings.Contains(call, ".") {
+						// Check if it's a local function reference
+						for _, localFunc := range localFunctions {
+							if call == localFunc {
+								resolvedCalls = append(resolvedCalls, packageName+"."+call)
+								break
+							}
+						}
+					} else {
+						resolvedCalls = append(resolvedCalls, call)
+					}
+				}
+				funcInfo.Calls = resolvedCalls
 			}
 			funcs = append(funcs, funcInfo)
 		}

cmd/server/main.go

@@ -439,19 +439,52 @@ func findFunctions(filePath, absPath, module string) ([]analyzer.FunctionInfo, e
 	if err != nil {
 		return nil, err
 	}
+
+	// Find package name
 	var packageName string
 	for _, line := range lines {
 		if strings.HasPrefix(line, "package ") {
 			packageName = strings.TrimSpace(strings.TrimPrefix(line, "package "))
 			break
 		}
 	}
+
+	// Collect all function names in this file for reference resolution
+	var localFunctions []string
+	for _, line := range lines {
+		if matches := re.FindStringSubmatch(line); matches != nil {
+			localFunctions = append(localFunctions, matches[1])
+		}
+	}
+
 	for i, line := range lines {
 		if matches := re.FindStringSubmatch(line); matches != nil {
-			fi := analyzer.FunctionInfo{Name: packageName + "." + matches[1], Line: i + 1, FilePath: relPath}
+			fi := analyzer.FunctionInfo{
+				Name:     packageName + "." + matches[1],
+				Line:     i + 1,
+				FilePath: relPath,
+			}
+			// Find function body
 			start, end := analyzer.FindFunctionBody(lines, i)
 			if start != -1 && end != -1 && start+1 < end && end < len(lines) {
-				fi.Calls = analyzer.FindCalls(lines[start+1 : end])
+				calls := analyzer.FindCalls(lines[start+1 : end])
+
+				// Resolve local function references by adding package prefix
+				var resolvedCalls []string
+				for _, call := range calls {
+					if !strings.Contains(call, ".") {
+						// Check if it's a local function reference
+						for _, localFunc := range localFunctions {
+							if call == localFunc {
+								resolvedCalls = append(resolvedCalls, packageName+"."+call)
+								break
+							}
+						}
+					} else {
+						resolvedCalls = append(resolvedCalls, call)
+					}
+				}
+				fi.Calls = resolvedCalls
 			}
 			funcs = append(funcs, fi)
 		}