goscript 0.0.49 → 0.0.51

package/compiler/analysis_test.go

@@ -160,6 +160,7 @@ func parseAndAnalyze(t *testing.T, code string) (*Analysis, map[string]types.Obj
 			Defs:  make(map[*ast.Ident]types.Object),
 			Uses:  make(map[*ast.Ident]types.Object),
 		},
+		Fset: fset,
 	}
 
 	// Type check the package
@@ -170,10 +171,8 @@ func parseAndAnalyze(t *testing.T, code string) (*Analysis, map[string]types.Obj
 	}
 	pkg.Types = typePkg
 
-	// Run analysis
-	analysis := NewAnalysis()
-	cmap := ast.NewCommentMap(fset, file, file.Comments)
-	AnalyzeFile(file, pkg, analysis, cmap)
+	// Run package-level analysis
+	analysis := AnalyzePackageFiles(pkg, nil)
 
 	// Collect variable objects
 	objects := make(map[string]types.Object)
@@ -269,18 +268,16 @@ func TestWrapperTypeDetection(t *testing.T) {
 		t.Fatalf("Package has errors: %v", pkg.Errors[0])
 	}
 
-	// Run analysis on the first file
+	// Run package-level analysis
 	if len(pkg.Syntax) == 0 {
 		t.Fatal("No syntax files found")
 	}
 
-	analysis := NewAnalysis()
-	cmap := ast.NewCommentMap(pkg.Fset, pkg.Syntax[0], pkg.Syntax[0].Comments)
-	AnalyzeFile(pkg.Syntax[0], pkg, analysis, cmap)
+	analysis := AnalyzePackageFiles(pkg, nil)
 
-	// Verify the WrapperTypes map was initialized
-	if analysis.WrapperTypes == nil {
-		t.Error("WrapperTypes map was not initialized")
+	// Verify the NamedBasicTypes map was initialized
+	if analysis.NamedBasicTypes == nil {
+		t.Error("NamedBasicTypes map was not initialized")
 	}
 
 	// Test some type lookups to verify wrapper type detection works
@@ -290,7 +287,7 @@ func TestWrapperTypeDetection(t *testing.T) {
 	// Check if MyMode is detected as a wrapper type
 	if obj := scope.Lookup("MyMode"); obj != nil {
 		if typeName, ok := obj.(*types.TypeName); ok {
-			isWrapper := analysis.IsWrapperType(typeName.Type())
+			isWrapper := analysis.IsNamedBasicType(typeName.Type())
 			if !isWrapper {
 				t.Errorf("MyMode should be detected as wrapper type, got %v", isWrapper)
 			}
@@ -301,7 +298,7 @@ func TestWrapperTypeDetection(t *testing.T) {
 	// Test that regular struct types are not detected as wrapper types
 	if obj := scope.Lookup("MyDir"); obj != nil {
 		if typeName, ok := obj.(*types.TypeName); ok {
-			isWrapper := analysis.IsWrapperType(typeName.Type())
+			isWrapper := analysis.IsNamedBasicType(typeName.Type())
 			if isWrapper {
 				t.Errorf("MyDir should not be detected as wrapper type, got %v", isWrapper)
 			}
@@ -309,13 +306,13 @@ func TestWrapperTypeDetection(t *testing.T) {
 		}
 	}
 
-	t.Logf("Analysis completed successfully with %d wrapper types tracked", len(analysis.WrapperTypes))
+	t.Logf("Analysis completed successfully with %d named basic types tracked", len(analysis.NamedBasicTypes))
 }
 
 // TestDiscoverGsPackages verifies that the discoverEmbeddedGsPackages function
 // can find packages in the embedded gs/ directory
 func TestDiscoverGsPackages(t *testing.T) {
-	analysis := NewAnalysis()
+	analysis := NewAnalysis(nil)
 
 	// Test package discovery using the embedded filesystem
 	packages := analysis.discoverEmbeddedGsPackages()

package/compiler/compiler.go

@@ -72,7 +72,11 @@ func NewCompiler(conf *Config, le *logrus.Entry, opts *packages.Config) (*Compil
 		packages.NeedTypesInfo |
 		packages.NeedTypesSizes
 
-	return &Compiler{config: *conf, le: le, opts: *opts}, nil
+	return &Compiler{
+		config: *conf,
+		le:     le,
+		opts:   *opts,
+	}, nil
 }
 
 // CompilationResult contains information about what was compiled
@@ -141,7 +145,6 @@ func (c *Compiler) CompilePackages(ctx context.Context, patterns ...string) (*Co
 			// Check if this package has a handwritten equivalent
 			if hasHandwrittenEquivalent(pkg.PkgPath) {
 				// Add this package but don't visit its dependencies
-				// c.le.Debugf("Skipping dependencies of handwritten package: %s", pkg.PkgPath)
 				return
 			}
 
@@ -149,13 +152,11 @@ func (c *Compiler) CompilePackages(ctx context.Context, patterns ...string) (*Co
 			for _, imp := range pkg.Imports {
 				// Skip protobuf-go-lite packages and their dependencies
 				if isProtobufGoLitePackage(imp.PkgPath) {
-					// c.le.Debugf("Skipping protobuf-go-lite package: %s", imp.PkgPath)
 					continue
 				}
 
 				// Skip packages that are only used by .pb.go files
 				if isPackageOnlyUsedByProtobufFiles(pkg, imp.PkgPath) {
-					// c.le.Debugf("Skipping package only used by .pb.go files: %s", imp.PkgPath)
 					continue
 				}
 
@@ -168,28 +169,39 @@ func (c *Compiler) CompilePackages(ctx context.Context, patterns ...string) (*Co
 			visit(pkg)
 		}
 
-		// Replace pkgs with all packages
-		pkgs = allPkgs
+		// Now we have collected all dependencies, but they only have minimal information.
+		// We need to reload them with complete type information for compilation.
+		// Collect all package paths from the dependency graph
+		var pkgPaths []string
+		pkgPathSet := make(map[string]bool) // Use set to avoid duplicates
 
-		/*
-			// Now load all packages with full mode to get complete type information
-			var pkgPaths []string
-			for _, pkg := range pkgs {
-				if pkg.PkgPath != "" {
-					pkgPaths = append(pkgPaths, pkg.PkgPath)
-				}
+		for _, pkg := range allPkgs {
+			if pkg.PkgPath != "" && !pkgPathSet[pkg.PkgPath] {
+				pkgPaths = append(pkgPaths, pkg.PkgPath)
+				pkgPathSet[pkg.PkgPath] = true
 			}
+		}
+
+		// Reload all collected packages with complete type information
+		if len(pkgPaths) > 0 {
+			c.le.Debugf("Reloading %d packages with complete type information", len(pkgPaths))
 
-			// Reload all packages with full mode
-			// TODO: Can we get rid of this? This would be very slow!
 			fullOpts := c.opts
 			fullOpts.Context = ctx
+			// Use LoadAllSyntax to get complete type information, syntax trees, and type checking
 			fullOpts.Mode = packages.LoadAllSyntax
-			pkgs, err = packages.Load(&fullOpts, pkgPaths...)
+
+			reloadedPkgs, err := packages.Load(&fullOpts, pkgPaths...)
 			if err != nil {
-				return fmt.Errorf("failed to reload packages with full mode: %w", err)
+				return nil, fmt.Errorf("failed to reload packages with complete type information: %w", err)
 			}
-		*/
+
+			// Replace the minimal packages with the fully loaded ones
+			pkgs = reloadedPkgs
+		} else {
+			// No packages to reload, use the original set
+			pkgs = allPkgs
+		}
 	}
 
 	// If DisableEmitBuiltin is false, we need to copy the builtin package to the output directory
@@ -206,6 +218,12 @@ func (c *Compiler) CompilePackages(ctx context.Context, patterns ...string) (*Co
 	// Track which gs packages have been processed to avoid duplicates
 	processedGsPackages := make(map[string]bool)
 
+	// Create a map of all loaded packages for dependency analysis
+	allPackages := make(map[string]*packages.Package)
+	for _, pkg := range pkgs {
+		allPackages[pkg.PkgPath] = pkg
+	}
+
 	// Compile all packages
 	for _, pkg := range pkgs {
 		// Check if the package has a handwritten equivalent
@@ -237,7 +255,7 @@ func (c *Compiler) CompilePackages(ctx context.Context, patterns ...string) (*Co
 			continue
 		}
 
-		pkgCompiler, err := NewPackageCompiler(c.le, &c.config, pkg)
+		pkgCompiler, err := NewPackageCompiler(c.le, &c.config, pkg, allPackages)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create package compiler for %s: %w", pkg.PkgPath, err)
 		}
@@ -262,6 +280,7 @@ type PackageCompiler struct {
 	compilerConf *Config
 	outputPath   string
 	pkg          *packages.Package
+	allPackages  map[string]*packages.Package
 }
 
 // NewPackageCompiler creates a new `PackageCompiler` for a given Go package.
@@ -272,12 +291,14 @@ func NewPackageCompiler(
 	le *logrus.Entry,
 	compilerConf *Config,
 	pkg *packages.Package,
+	allPackages map[string]*packages.Package,
 ) (*PackageCompiler, error) {
 	res := &PackageCompiler{
 		le:           le,
 		pkg:          pkg,
 		compilerConf: compilerConf,
 		outputPath:   ComputeModulePath(compilerConf.OutputPath, pkg.PkgPath),
+		allPackages:  allPackages,
 	}
 
 	return res, nil
@@ -304,7 +325,10 @@ func (c *PackageCompiler) Compile(ctx context.Context) error {
 	}
 
 	// Perform package-level analysis for auto-imports
-	packageAnalysis := AnalyzePackage(c.pkg)
+	packageAnalysis := AnalyzePackageImports(c.pkg)
+
+	// Perform comprehensive package-level analysis for code generation
+	analysis := AnalyzePackageFiles(c.pkg, c.allPackages)
 
 	// Track all compiled files for later generating the index.ts
 	compiledFiles := make([]string, 0, len(c.pkg.CompiledGoFiles))
@@ -342,7 +366,7 @@ func (c *PackageCompiler) Compile(ctx context.Context) error {
 
 		// log just the filename
 		c.le.Debugf("GS: %s", filepath.Base(fileName))
-		if err := c.CompileFile(ctx, fileName, f, packageAnalysis); err != nil {
+		if err := c.CompileFile(ctx, fileName, f, analysis, packageAnalysis); err != nil {
 			return err
 		}
 
@@ -392,6 +416,7 @@ func (c *PackageCompiler) generateIndexFile(compiledFiles []string) error {
 		// Find which symbols this file exports
 		var valueSymbols []string
 		var typeSymbols []string
+		var structSymbols []string // New: Track structs separately
 
 		// Find the corresponding syntax file
 		for i, syntax := range c.pkg.Syntax {
@@ -408,21 +433,27 @@ func (c *PackageCompiler) generateIndexFile(compiledFiles []string) error {
 				switch d := decl.(type) {
 				case *ast.FuncDecl:
 					if d.Recv == nil && d.Name.IsExported() {
-						valueSymbols = append(valueSymbols, d.Name.Name)
+						valueSymbols = append(valueSymbols, sanitizeIdentifier(d.Name.Name))
 					}
 				case *ast.GenDecl:
 					for _, spec := range d.Specs {
 						switch s := spec.(type) {
 						case *ast.TypeSpec:
 							if s.Name.IsExported() {
-								// All type declarations (interfaces, structs, type definitions, type aliases)
-								// become TypeScript types and must be exported with "export type"
-								typeSymbols = append(typeSymbols, s.Name.Name)
+								// Check if this is a struct type
+								if _, isStruct := s.Type.(*ast.StructType); isStruct {
+									// Structs become TypeScript classes and need both type and value exports
+									structSymbols = append(structSymbols, sanitizeIdentifier(s.Name.Name))
+								} else {
+									// Other type declarations (interfaces, type definitions, type aliases)
+									// become TypeScript types and must be exported with "export type"
+									typeSymbols = append(typeSymbols, sanitizeIdentifier(s.Name.Name))
+								}
 							}
 						case *ast.ValueSpec:
 							for _, name := range s.Names {
 								if name.IsExported() {
-									valueSymbols = append(valueSymbols, name.Name)
+									valueSymbols = append(valueSymbols, sanitizeIdentifier(name.Name))
 								}
 							}
 						}
@@ -442,6 +473,17 @@ func (c *PackageCompiler) generateIndexFile(compiledFiles []string) error {
 			}
 		}
 
+		// Write struct exports (both as types and values)
+		if len(structSymbols) > 0 {
+			sort.Strings(structSymbols)
+			// Export classes as values (which makes them available as both types and values in TypeScript)
+			exportLine := fmt.Sprintf("export { %s } from \"./%s.js\"\n",
+				strings.Join(structSymbols, ", "), fileName)
+			if _, err := indexFile.WriteString(exportLine); err != nil {
+				return err
+			}
+		}
+
 		if len(typeSymbols) > 0 {
 			sort.Strings(typeSymbols)
 			exportLine := fmt.Sprintf("export type { %s } from \"./%s.js\"\n",
@@ -456,21 +498,11 @@ func (c *PackageCompiler) generateIndexFile(compiledFiles []string) error {
 }
 
 // CompileFile handles the compilation of a single Go source file to TypeScript.
-// It first performs a pre-compilation analysis of the file using `AnalyzeFile`
-// to gather information necessary for accurate TypeScript generation (e.g.,
-// about varRefing, async functions, defer statements).
+// It uses the pre-computed package-level analysis for accurate TypeScript generation
+// (e.g., about varRefing, async functions, defer statements, receiver usage across files).
 // Then, it creates a `FileCompiler` instance for the file and invokes its
 // `Compile` method to generate the TypeScript code.
-func (p *PackageCompiler) CompileFile(ctx context.Context, name string, syntax *ast.File, packageAnalysis *PackageAnalysis) error {
-	// Create a new analysis instance for per-file data
-	analysis := NewAnalysis()
-
-	// Create comment map for the file
-	cmap := ast.NewCommentMap(p.pkg.Fset, syntax, syntax.Comments)
-
-	// Analyze the file before compiling
-	AnalyzeFile(syntax, p.pkg, analysis, cmap)
-
+func (p *PackageCompiler) CompileFile(ctx context.Context, name string, syntax *ast.File, analysis *Analysis, packageAnalysis *PackageAnalysis) error {
 	fileCompiler, err := NewFileCompiler(p.compilerConf, p.pkg, syntax, name, analysis, packageAnalysis)
 	if err != nil {
 		return err
@@ -644,6 +676,10 @@ type GoToTSCompiler struct {
 	pkg *packages.Package
 
 	analysis *Analysis
+
+	// awaitedCalls tracks which call expressions have already been processed
+	// to avoid adding double await keywords
+	awaitedCalls map[*ast.CallExpr]bool
 }
 
 // It initializes the compiler with a `TSCodeWriter` for output,

package/compiler/decl.go

@@ -3,6 +3,7 @@ package compiler
 import (
 	"fmt"
 	"go/ast"
+	"go/types"
 )
 
 // WriteDecls iterates through a slice of Go top-level declarations (`ast.Decl`)
@@ -193,6 +194,27 @@ func (c *GoToTSCompiler) writeMethodSignature(decl *ast.FuncDecl) (bool, error)
 	var isAsync bool
 	if obj := c.pkg.TypesInfo.Defs[decl.Name]; obj != nil {
 		isAsync = c.analysis.IsAsyncFunc(obj)
+
+		// Check if this method must be async due to interface constraints
+		if !isAsync && decl.Recv != nil && len(decl.Recv.List) > 0 {
+			// Get the receiver type
+			receiverType := decl.Recv.List[0].Type
+			if starExpr, ok := receiverType.(*ast.StarExpr); ok {
+				receiverType = starExpr.X
+			}
+
+			if ident, ok := receiverType.(*ast.Ident); ok {
+				// Get the named type for this receiver
+				if receiverObj := c.pkg.TypesInfo.Uses[ident]; receiverObj != nil {
+					if namedType, ok := receiverObj.Type().(*types.Named); ok {
+						// Check if this method must be async due to interface constraints
+						if c.analysis.MustBeAsyncDueToInterface(namedType, decl.Name.Name) {
+							isAsync = true
+						}
+					}
+				}
+			}
+		}
 	}
 
 	// Methods are typically public in the TS output

package/compiler/expr-call-async.go

@@ -2,74 +2,273 @@ package compiler
 
 import (
 	"go/ast"
+	"go/token"
 	"go/types"
 	"strings"
 )
 
-// writeAsyncCall writes the await prefix for async function calls
-func (c *GoToTSCompiler) writeAsyncCall(exp *ast.CallExpr, funIdent *ast.Ident) bool {
-	if funIdent == nil {
-		return false
+// writeAsyncCallIfNeeded writes the await prefix for async function or method calls
+// Returns true if await was written, false otherwise
+func (c *GoToTSCompiler) writeAsyncCallIfNeeded(exp *ast.CallExpr) bool {
+	// Track if we've already processed this call expression to avoid double await
+	if c.awaitedCalls == nil {
+		c.awaitedCalls = make(map[*ast.CallExpr]bool)
 	}
 
-	// Check if this is an async function call
-	if obj := c.pkg.TypesInfo.Uses[funIdent]; obj != nil && c.analysis.IsAsyncFunc(obj) {
-		c.tsw.WriteLiterally("await ")
-		return true
+	if c.awaitedCalls[exp] {
+		return false // Already processed this call
 	}
+	switch fun := exp.Fun.(type) {
+	case *ast.Ident:
+		// Function call (e.g., func())
+		if obj := c.pkg.TypesInfo.Uses[fun]; obj != nil && c.analysis.IsAsyncFunc(obj) {
+			c.awaitedCalls[exp] = true
+			c.tsw.WriteLiterally("await ")
+			return true
+		}
+		return false
 
-	return false
-}
+	case *ast.SelectorExpr:
+		// Method call (e.g., obj.method())
+		var ident *ast.Ident
+		var identOk bool
+
+		// Handle both direct identifiers and pointer dereferences
+		switch x := fun.X.(type) {
+		case *ast.Ident:
+			ident, identOk = x, true
+		case *ast.StarExpr:
+			// Handle pointer dereference: (*p).method() or p.method() where p is a pointer
+			if id, isIdent := x.X.(*ast.Ident); isIdent {
+				ident, identOk = id, true
+			}
+		default:
+			return false
+		}
+
+		if !identOk {
+			return false
+		}
+
+		// Get the type of the receiver
+		obj := c.pkg.TypesInfo.Uses[ident]
+		if obj == nil {
+			return false
+		}
+
+		varObj, ok := obj.(*types.Var)
+		if !ok {
+			return false
+		}
+
+		// Get the type name and package
+		var namedType *types.Named
+		var namedTypeOk bool
+
+		// Handle both direct named types and pointer to named types
+		switch t := varObj.Type().(type) {
+		case *types.Named:
+			namedType, namedTypeOk = t, true
+		case *types.Pointer:
+			if nt, isNamed := t.Elem().(*types.Named); isNamed {
+				namedType, namedTypeOk = nt, true
+			}
+		}
+
+		if !namedTypeOk {
+			return false
+		}
+
+		typeName := namedType.Obj().Name()
+		methodName := fun.Sel.Name
 
-// writeAsyncMethodCall writes the await prefix for async method calls
-func (c *GoToTSCompiler) writeAsyncMethodCall(exp *ast.CallExpr) bool {
-	selExpr, ok := exp.Fun.(*ast.SelectorExpr)
-	if !ok {
+		// Check if the type is from an imported package
+		typePkg := namedType.Obj().Pkg()
+		if typePkg != nil && typePkg != c.pkg.Types {
+			// Use the full package path from the type information (not just the package name)
+			pkgPath := typePkg.Path()
+
+			// Check if this method is async based on metadata
+			if c.analysis.IsMethodAsync(pkgPath, typeName, methodName) {
+				c.awaitedCalls[exp] = true
+				c.tsw.WriteLiterally("await ")
+				return true
+			}
+		} else {
+			// For local types, check if the method contains async operations at runtime
+			// to avoid circular dependencies during analysis
+			if c.isLocalMethodAsync(namedType, methodName) {
+				c.awaitedCalls[exp] = true
+				c.tsw.WriteLiterally("await ")
+				return true
+			}
+		}
 		return false
-	}
 
-	// Check if this is a method call on a variable (e.g., mu.Lock())
-	ident, ok := selExpr.X.(*ast.Ident)
-	if !ok {
+	default:
 		return false
 	}
+}
 
-	// Get the type of the receiver
-	obj := c.pkg.TypesInfo.Uses[ident]
-	if obj == nil {
-		return false
+// isLocalMethodAsync checks if a local method contains async operations by inspecting its body
+func (c *GoToTSCompiler) isLocalMethodAsync(namedType *types.Named, methodName string) bool {
+	// Find the method in the named type
+	for i := 0; i < namedType.NumMethods(); i++ {
+		method := namedType.Method(i)
+		if method.Name() == methodName {
+			// Find the method declaration in the AST
+			methodDecl := c.findMethodDecl(namedType, methodName)
+			if methodDecl != nil && methodDecl.Body != nil {
+				// Check if the method body contains async operations
+				return c.containsAsyncOperationsRuntime(methodDecl.Body)
+			}
+			break
+		}
+	}
+	return false
+}
+
+// findMethodDecl finds the AST declaration for a method on a named type
+func (c *GoToTSCompiler) findMethodDecl(namedType *types.Named, methodName string) *ast.FuncDecl {
+	// Search through all files in the package for the method declaration
+	for _, file := range c.pkg.Syntax {
+		for _, decl := range file.Decls {
+			if funcDecl, ok := decl.(*ast.FuncDecl); ok {
+				// Check if this is a method with the right name
+				if funcDecl.Name.Name == methodName && funcDecl.Recv != nil {
+					// Check if the receiver type matches
+					if c.isReceiverOfType(funcDecl, namedType) {
+						return funcDecl
+					}
+				}
+			}
+		}
 	}
+	return nil
+}
 
-	varObj, ok := obj.(*types.Var)
-	if !ok {
+// isReceiverOfType checks if a function declaration is a method of the given named type
+func (c *GoToTSCompiler) isReceiverOfType(funcDecl *ast.FuncDecl, namedType *types.Named) bool {
+	if funcDecl.Recv == nil || len(funcDecl.Recv.List) == 0 {
 		return false
 	}
 
-	// Get the type name and package
-	namedType, ok := varObj.Type().(*types.Named)
-	if !ok {
+	recvField := funcDecl.Recv.List[0]
+	if len(recvField.Names) == 0 {
 		return false
 	}
 
-	typeName := namedType.Obj().Name()
-	methodName := selExpr.Sel.Name
+	recvIdent := recvField.Names[0]
+	if obj := c.pkg.TypesInfo.Defs[recvIdent]; obj != nil {
+		if varObj, ok := obj.(*types.Var); ok {
+			// Handle both direct named types and pointer to named types
+			var receiverNamedType *types.Named
+			switch t := varObj.Type().(type) {
+			case *types.Named:
+				receiverNamedType = t
+			case *types.Pointer:
+				if nt, isNamed := t.Elem().(*types.Named); isNamed {
+					receiverNamedType = nt
+				}
+			}
 
-	// Check if the type is from an imported package
-	typePkg := namedType.Obj().Pkg()
-	if typePkg == nil || typePkg == c.pkg.Types {
-		return false
+			if receiverNamedType != nil {
+				// For generic types, compare the origin types and names
+				// since instantiated generics have different type objects
+				targetObj := namedType.Obj()
+				receiverObj := receiverNamedType.Obj()
+
+				// Compare by name and package path
+				return targetObj.Name() == receiverObj.Name() &&
+					targetObj.Pkg() == receiverObj.Pkg()
+			}
+		}
 	}
 
-	// Use the actual package name from the type information
-	pkgName := typePkg.Name()
+	return false
+}
+
+// containsAsyncOperationsRuntime checks if a node contains async operations at runtime
+func (c *GoToTSCompiler) containsAsyncOperationsRuntime(node ast.Node) bool {
+	var hasAsync bool
+
+	ast.Inspect(node, func(n ast.Node) bool {
+		if n == nil {
+			return false
+		}
+
+		switch s := n.(type) {
+		case *ast.SendStmt:
+			// Channel send operation (ch <- value)
+			hasAsync = true
+			return false
+
+		case *ast.UnaryExpr:
+			// Channel receive operation (<-ch)
+			if s.Op == token.ARROW {
+				hasAsync = true
+				return false
+			}
+
+		case *ast.SelectStmt:
+			// Select statement with channel operations
+			hasAsync = true
+			return false
+
+		case *ast.CallExpr:
+			// Check if we're calling a function known to be async
+			if funcIdent, ok := s.Fun.(*ast.Ident); ok {
+				// Get the object for this function call
+				if obj := c.pkg.TypesInfo.Uses[funcIdent]; obj != nil && c.analysis.IsAsyncFunc(obj) {
+					hasAsync = true
+					return false
+				}
+			}
+
+			// Check for method calls on imported types (similar to analysis.go logic)
+			if selExpr, ok := s.Fun.(*ast.SelectorExpr); ok {
+				if ident, ok := selExpr.X.(*ast.Ident); ok {
+					if obj := c.pkg.TypesInfo.Uses[ident]; obj != nil {
+						if varObj, ok := obj.(*types.Var); ok {
+							var namedType *types.Named
+							switch t := varObj.Type().(type) {
+							case *types.Named:
+								namedType = t
+							case *types.Pointer:
+								if nt, isNamed := t.Elem().(*types.Named); isNamed {
+									namedType = nt
+								}
+							}
+
+							if namedType != nil {
+								typeName := namedType.Obj().Name()
+								methodName := selExpr.Sel.Name
+
+								// Check if the type is from an imported package
+								if typePkg := namedType.Obj().Pkg(); typePkg != nil && typePkg != c.pkg.Types {
+									pkgPath := typePkg.Path()
+									if c.analysis.IsMethodAsync(pkgPath, typeName, methodName) {
+										hasAsync = true
+										return false
+									}
+								} else {
+									// For local types, recursively check
+									if c.isLocalMethodAsync(namedType, methodName) {
+										hasAsync = true
+										return false
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
 
-	// Check if this method is async based on metadata
-	if c.analysis.IsMethodAsync(pkgName, typeName, methodName) {
-		c.tsw.WriteLiterally("await ")
 		return true
-	}
+	})
 
-	return false
+	return hasAsync
 }
 
 // addNonNullAssertion adds ! for function calls that might return null