fix: Universe LocateInPackage should use source dir checking

go.mod

@@ -3,7 +3,7 @@ module github.com/octohelm/gengo
 go 1.23.2
 
 require (
-	github.com/go-courier/logr v0.3.0
+	github.com/go-courier/logr v0.3.1
 	github.com/octohelm/x v0.0.0-20241011014327-0fcf864c84d6
 	github.com/onsi/gomega v1.34.2
 	golang.org/x/sync v0.8.0
@@ -13,10 +13,7 @@ require (
 
 require (
 	github.com/google/go-cmp v0.6.0 // indirect
-	github.com/kr/text v0.2.0 // indirect
-	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	golang.org/x/mod v0.21.0 // indirect
 	golang.org/x/net v0.30.0 // indirect
-	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,6 +1,5 @@
-github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
-github.com/go-courier/logr v0.3.0 h1:0VEQB1b53EmYQ+ZehrIgD8l2IO+WX7TY+CqzlykIFmo=
-github.com/go-courier/logr v0.3.0/go.mod h1:OI7f/JCFZ1ZMD5qG3bIJr5WMNnGzd24+II1D9D9w5x4=
+github.com/go-courier/logr v0.3.1 h1:RcHM7qpO8OpuV+zFvJMXtJEspTnnYmT6uGiAomwb8X8=
+github.com/go-courier/logr v0.3.1/go.mod h1:NQWi+TSv0rS1RfyWHv7MNEI5cNy9NR6k1n8R24uHgdY=
 github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
 github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
@@ -9,12 +8,6 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/pprof v0.0.0-20240827171923-fa2c70bbbfe5 h1:5iH8iuqE5apketRbSFBy+X1V0o+l+8NF1avt4HWl7cA=
 github.com/google/pprof v0.0.0-20240827171923-fa2c70bbbfe5/go.mod h1:vavhavw2zAxS5dIdcRluK6cSGGPlZynqzFM8NdvU144=
-github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
-github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
-github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
-github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/octohelm/x v0.0.0-20241011014327-0fcf864c84d6 h1:4royPn66/B49ftj5Jh7ZWuJ8A2B/2l0UTQ6EFLf29Vk=
 github.com/octohelm/x v0.0.0-20241011014327-0fcf864c84d6/go.mod h1:c8k4TZNwZXVCclWxFU8dm67PQ98APzKW3f5JJykQ2uo=
 github.com/onsi/ginkgo/v2 v2.20.1 h1:YlVIbqct+ZmnEph770q9Q7NVAz4wwIiVNahee6JyUzo=
@@ -33,8 +26,7 @@ golang.org/x/text v0.19.0 h1:kTxAhCbGbxhK0IwgSKiMO5awPoDQ0RpfiVYBfK860YM=
 golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
 golang.org/x/tools v0.26.0 h1:v/60pFQmzmT9ExmjDv2gGIfi3OqfKoEP6I5+umXlbnQ=
 golang.org/x/tools v0.26.0/go.mod h1:TPVVj70c7JJ3WCazhD8OdXcZg/og+b9+tH/KxylGwH0=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=
-gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/gengo/context.go

@@ -30,7 +30,7 @@ func NewContext(args *GeneratorArgs) (Executor, error) {
 	c := &context{
 		pkgPaths: pkgPaths,
 		args:     args,
-		universe: u,
+		universe: *u,
 	}
 	return c, nil
 }
@@ -47,9 +47,9 @@ type context struct {
 	pkgPaths map[string]bool
 	args     *GeneratorArgs
 
-	universe gengotypes.Universe
 	pkgTags  map[string][]string
 	pkg      gengotypes.Package
+	universe gengotypes.Universe
 	genfile  *genfile
 	l        logr.Logger
 }
@@ -84,8 +84,8 @@ func (c *context) pkgExecute(ctx corecontext.Context, pkg string, generators ...
 	ctx, l := logr.FromContext(ctx).Start(ctx, pkg)
 	defer l.End()
 
-	p, ok := c.universe[pkg]
-	if !ok {
+	p := c.universe.Package(pkg)
+	if p == nil {
 		return fmt.Errorf("invalid pkg `%s`", pkg)
 	}
 

pkg/types/load.go

@@ -3,6 +3,7 @@ package types
 import (
 	"fmt"
 	"go/token"
+	"path/filepath"
 
 	"golang.org/x/tools/go/packages"
 )
@@ -15,8 +16,11 @@ const (
 	LoadAllSyntax = LoadSyntax | packages.NeedDeps | packages.NeedModule
 )
 
-func Load(patterns []string) (Universe, map[string]bool, error) {
+func Load(patterns []string) (*Universe, map[string]bool, error) {
+	fset := token.NewFileSet()
+
 	c := &packages.Config{
+		Fset: fset,
 		Mode: LoadAllSyntax,
 	}
 
@@ -25,12 +29,14 @@ func Load(patterns []string) (Universe, map[string]bool, error) {
 		return nil, nil, err
 	}
 
-	u := Universe{}
+	u := &Universe{
+		fset: fset,
+		pkgs: map[string]Package{},
+	}
 	rootPkgPaths := map[string]bool{}
 	pkgPaths := map[string]bool{}
 
 	var register func(p *packages.Package)
-
 	register = func(p *packages.Package) {
 		if len(p.Errors) > 0 {
 			for i := range p.Errors {
@@ -44,20 +50,19 @@ func Load(patterns []string) (Universe, map[string]bool, error) {
 		for k := range p.Imports {
 			importedPkg := p.Imports[k]
 
-			if _, ok := u[importedPkg.PkgPath]; !ok {
+			if _, ok := u.pkgs[importedPkg.PkgPath]; !ok {
 				register(importedPkg)
 			}
 		}
 
-		u[p.PkgPath] = pkg
+		u.pkgs[p.PkgPath] = pkg
 
 		for rootPkgPath := range rootPkgPaths {
 			// when is sub pkg of root pkg
 			if p.Module != nil && rootPkgPath == p.Module.Path {
 				pkgPaths[p.PkgPath] = true
 			}
 		}
-
 	}
 
 	for i := range pkgs {
@@ -82,19 +87,23 @@ func Load(patterns []string) (Universe, map[string]bool, error) {
 	return u, pkgPaths, nil
 }
 
-type Universe map[string]Package
+type Universe struct {
+	fset *token.FileSet
+	pkgs map[string]Package
+}
 
-func (u Universe) Package(pkgPath string) Package {
-	return u[pkgPath]
+func (u *Universe) Package(pkgPath string) Package {
+	v, _ := u.pkgs[pkgPath]
+	return v
 }
 
-func (u Universe) LocateInPackage(pos token.Pos) Package {
-	for i := range u {
-		p := u[i]
-		for _, f := range p.Files() {
-			if f.Pos() >= pos && pos <= f.End() {
-				return p
-			}
+func (u *Universe) LocateInPackage(pos token.Pos) Package {
+	pp := u.fset.Position(pos)
+	dir := filepath.Dir(pp.Filename)
+
+	for _, p := range u.pkgs {
+		if dir == p.SourceDir() {
+			return p
 		}
 	}
 	return nil

pkg/types/package.go

@@ -1,14 +1,14 @@
 package types
 
 import (
+	"github.com/octohelm/x/ptr"
 	"go/ast"
 	"go/token"
 	"go/types"
+	"golang.org/x/tools/go/packages"
 	"path/filepath"
 	"strings"
 	"sync"
-
-	"golang.org/x/tools/go/packages"
 )
 
 type Module struct {
@@ -62,7 +62,7 @@ type ModInfo struct {
 	Require map[string]string
 }
 
-func newPkg(pkg *packages.Package, u Universe) Package {
+func newPkg(pkg *packages.Package, u *Universe) Package {
 	p := &pkgInfo{
 		u: u,
 
@@ -222,7 +222,7 @@ func newPkg(pkg *packages.Package, u Universe) Package {
 }
 
 type pkgInfo struct {
-	u       Universe
+	u       *Universe
 	Package *packages.Package
 
 	imports map[string]Package
@@ -240,13 +240,31 @@ type pkgInfo struct {
 
 	funcResults         sync.Map
 	funcResultResolvers sync.Map
+
+	sourceDir *string
 }
 
 func (pi *pkgInfo) SourceDir() string {
-	if pi.Package.PkgPath == pi.Module().Path {
-		return pi.Module().Dir
+	if pi.sourceDir != nil {
+		return *pi.sourceDir
 	}
-	return filepath.Join(pi.Module().Dir, pi.Package.PkgPath[len(pi.Module().Path):])
+
+	sourceDir := func() string {
+		if pi == nil || pi.Package == nil {
+			return ""
+		}
+		if pi.Module() == nil {
+			return ""
+		}
+		if pi.Package.PkgPath == pi.Module().Path {
+			return pi.Module().Dir
+		}
+		return filepath.Join(pi.Module().Dir, pi.Package.PkgPath[len(pi.Module().Path):])
+	}()
+
+	pi.sourceDir = ptr.Ptr(sourceDir)
+
+	return sourceDir
 }
 
 func (pi *pkgInfo) Pkg() *types.Package {