github.com/djui/moq@v0.3.3/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  
    86  	case *types.Array:
    87  		m.populateImports(t.Elem(), imports)
    88  
    89  	case *types.Slice:
    90  		m.populateImports(t.Elem(), imports)
    91  
    92  	case *types.Signature:
    93  		for i := 0; i < t.Params().Len(); i++ {
    94  			m.populateImports(t.Params().At(i).Type(), imports)
    95  		}
    96  		for i := 0; i < t.Results().Len(); i++ {
    97  			m.populateImports(t.Results().At(i).Type(), imports)
    98  		}
    99  
   100  	case *types.Map:
   101  		m.populateImports(t.Key(), imports)
   102  		m.populateImports(t.Elem(), imports)
   103  
   104  	case *types.Chan:
   105  		m.populateImports(t.Elem(), imports)
   106  
   107  	case *types.Pointer:
   108  		m.populateImports(t.Elem(), imports)
   109  
   110  	case *types.Struct: // anonymous struct
   111  		for i := 0; i < t.NumFields(); i++ {
   112  			m.populateImports(t.Field(i).Type(), imports)
   113  		}
   114  
   115  	case *types.Interface: // anonymous interface
   116  		for i := 0; i < t.NumExplicitMethods(); i++ {
   117  			m.populateImports(t.ExplicitMethod(i).Type(), imports)
   118  		}
   119  		for i := 0; i < t.NumEmbeddeds(); i++ {
   120  			m.populateImports(t.EmbeddedType(i), imports)
   121  		}
   122  	}
   123  }
   124  
   125  // resolveImportVarConflicts ensures that all the newly added imports do not
   126  // conflict with any of the existing vars.
   127  func (m MethodScope) resolveImportVarConflicts(imports map[string]*Package) {
   128  	// Ensure that all the newly added imports do not conflict with any of the
   129  	// existing vars.
   130  	for _, imprt := range imports {
   131  		if v, ok := m.searchVar(imprt.Qualifier()); ok {
   132  			v.Name += "MoqParam"
   133  		}
   134  	}
   135  }