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  }