github.com/moitias/moq@v0.0.0-20240223074357-5eb0f0ba4054/internal/registry/method_scope.go (about) 1 package registry 2 3 import ( 4 "go/types" 5 "strconv" 6 ) 7 8 // MethodScope is the sub-registry for allocating variables present in 9 // the method scope. 10 // 11 // It should be created using a registry instance. 12 type MethodScope struct { 13 registry *Registry 14 moqPkgPath string 15 16 vars []*Var 17 conflicted map[string]bool 18 } 19 20 // AddVar allocates a variable instance and adds it to the method scope. 21 // 22 // Variables names are generated if required and are ensured to be 23 // without conflict with other variables and imported packages. It also 24 // adds the relevant imports to the registry for each added variable. 25 func (m *MethodScope) AddVar(vr *types.Var, suffix string) *Var { 26 imports := make(map[string]*Package) 27 m.populateImports(vr.Type(), imports) 28 m.resolveImportVarConflicts(imports) 29 30 name := varName(vr, suffix) 31 // Ensure that the var name does not conflict with a package import. 32 if _, ok := m.registry.searchImport(name); ok { 33 name += "MoqParam" 34 } 35 if _, ok := m.searchVar(name); ok || m.conflicted[name] { 36 name = m.resolveVarNameConflict(name) 37 } 38 39 v := Var{ 40 vr: vr, 41 imports: imports, 42 moqPkgPath: m.moqPkgPath, 43 Name: name, 44 } 45 m.vars = append(m.vars, &v) 46 return &v 47 } 48 49 func (m *MethodScope) resolveVarNameConflict(suggested string) string { 50 for n := 1; ; n++ { 51 _, ok := m.searchVar(suggested + strconv.Itoa(n)) 52 if ok { 53 continue 54 } 55 56 if n == 1 { 57 conflict, _ := m.searchVar(suggested) 58 conflict.Name += "1" 59 m.conflicted[suggested] = true 60 n++ 61 } 62 return suggested + strconv.Itoa(n) 63 } 64 } 65 66 func (m MethodScope) searchVar(name string) (*Var, bool) { 67 for _, v := range m.vars { 68 if v.Name == name { 69 return v, true 70 } 71 } 72 73 return nil, false 74 } 75 76 // populateImports extracts all the package imports for a given type 77 // recursively. The imported packages by a single type can be more than 78 // one (ex: map[a.Type]b.Type). 79 func (m MethodScope) populateImports(t types.Type, imports map[string]*Package) { 80 switch t := t.(type) { 81 case *types.Named: 82 if pkg := t.Obj().Pkg(); pkg != nil { 83 imports[stripVendorPath(pkg.Path())] = m.registry.AddImport(pkg) 84 } 85 // The imports of a Type with a TypeList must be added to the imports list 86 // For example: Foo[otherpackage.Bar] , must have otherpackage imported 87 if targs := t.TypeArgs(); targs != nil { 88 for i := 0; i < targs.Len(); i++ { 89 m.populateImports(targs.At(i), imports) 90 } 91 } 92 93 case *types.Array: 94 m.populateImports(t.Elem(), imports) 95 96 case *types.Slice: 97 m.populateImports(t.Elem(), imports) 98 99 case *types.Signature: 100 for i := 0; i < t.Params().Len(); i++ { 101 m.populateImports(t.Params().At(i).Type(), imports) 102 } 103 for i := 0; i < t.Results().Len(); i++ { 104 m.populateImports(t.Results().At(i).Type(), imports) 105 } 106 107 case *types.Map: 108 m.populateImports(t.Key(), imports) 109 m.populateImports(t.Elem(), imports) 110 111 case *types.Chan: 112 m.populateImports(t.Elem(), imports) 113 114 case *types.Pointer: 115 m.populateImports(t.Elem(), imports) 116 117 case *types.Struct: // anonymous struct 118 for i := 0; i < t.NumFields(); i++ { 119 m.populateImports(t.Field(i).Type(), imports) 120 } 121 122 case *types.Interface: // anonymous interface 123 for i := 0; i < t.NumExplicitMethods(); i++ { 124 m.populateImports(t.ExplicitMethod(i).Type(), imports) 125 } 126 for i := 0; i < t.NumEmbeddeds(); i++ { 127 m.populateImports(t.EmbeddedType(i), imports) 128 } 129 } 130 } 131 132 // resolveImportVarConflicts ensures that all the newly added imports do not 133 // conflict with any of the existing vars. 134 func (m MethodScope) resolveImportVarConflicts(imports map[string]*Package) { 135 // Ensure that all the newly added imports do not conflict with any of the 136 // existing vars. 137 for _, imprt := range imports { 138 if v, ok := m.searchVar(imprt.Qualifier()); ok { 139 v.Name += "MoqParam" 140 } 141 } 142 }