moq/internal/registry/method_scope.go

package registry

import (
	"go/types"
	"strconv"
)

// MethodScope is the sub-registry for allocating variables present in
// the method scope.
//
// It should be created using a registry instance.
type MethodScope struct {
	registry   *Registry
	moqPkgPath string

	vars       []*Var
	conflicted map[string]bool
}

// AddVar allocates a variable instance and adds it to the method scope.
//
// Variables names are generated if required and are ensured to be
// without conflict with other variables and imported packages. It also
// adds the relevant imports to the registry for each added variable.
func (m *MethodScope) AddVar(vr *types.Var, suffix string) *Var {
	name := vr.Name()
	if name == "" || name == "_" {
		name = generateVarName(vr.Type())
	}

	name += suffix

	switch name {
	case "mock", "callInfo", "break", "default", "func", "interface", "select", "case", "defer", "go", "map", "struct",
		"chan", "else", "goto", "package", "switch", "const", "fallthrough", "if", "range", "type", "continue", "for",
		"import", "return", "var":
		name += "MoqParam"
	}

	if _, ok := m.searchVar(name); ok || m.conflicted[name] {
		return m.addDisambiguatedVar(vr, name)
	}

	return m.addVar(vr, name)
}

func (m *MethodScope) addDisambiguatedVar(vr *types.Var, suggested string) *Var {
	n := 1
	for {
		// Keep incrementing the suffix until we find a name which is unused.
		if _, ok := m.searchVar(suggested + strconv.Itoa(n)); !ok {
			break
		}
		n++
	}

	name := suggested + strconv.Itoa(n)
	if n == 1 {
		conflict, _ := m.searchVar(suggested)
		conflict.Name += "1"
		name = suggested + "2"
		m.conflicted[suggested] = true
	}

	return m.addVar(vr, name)
}

func (m *MethodScope) addVar(vr *types.Var, name string) *Var {
	imports := make(map[string]*Package)
	m.populateImports(vr.Type(), imports)

	v := Var{
		vr:         vr,
		imports:    imports,
		moqPkgPath: m.moqPkgPath,
		Name:       name,
	}
	m.vars = append(m.vars, &v)
	m.resolveImportVarConflicts(&v)
	return &v
}

func (m MethodScope) searchVar(name string) (*Var, bool) {
	for _, v := range m.vars {
		if v.Name == name {
			return v, true
		}
	}

	return nil, false
}

// populateImports extracts all the package imports for a given type
// recursively. The imported packages by a single type can be more than
// one (ex: map[a.Type]b.Type).
func (m MethodScope) populateImports(t types.Type, imports map[string]*Package) {
	switch t := t.(type) {
	case *types.Named:
		if pkg := t.Obj().Pkg(); pkg != nil {
			imports[stripVendorPath(pkg.Path())] = m.registry.AddImport(pkg)
		}

	case *types.Array:
		m.populateImports(t.Elem(), imports)

	case *types.Slice:
		m.populateImports(t.Elem(), imports)

	case *types.Signature:
		for i := 0; i < t.Params().Len(); i++ {
			m.populateImports(t.Params().At(i).Type(), imports)
		}
		for i := 0; i < t.Results().Len(); i++ {
			m.populateImports(t.Results().At(i).Type(), imports)
		}

	case *types.Map:
		m.populateImports(t.Key(), imports)
		m.populateImports(t.Elem(), imports)

	case *types.Chan:
		m.populateImports(t.Elem(), imports)

	case *types.Pointer:
		m.populateImports(t.Elem(), imports)

	case *types.Struct: // anonymous struct
		for i := 0; i < t.NumFields(); i++ {
			m.populateImports(t.Field(i).Type(), imports)
		}

	case *types.Interface: // anonymous interface
		for i := 0; i < t.NumExplicitMethods(); i++ {
			m.populateImports(t.ExplicitMethod(i).Type(), imports)
		}
		for i := 0; i < t.NumEmbeddeds(); i++ {
			m.populateImports(t.EmbeddedType(i), imports)
		}
	}
}

func (m MethodScope) resolveImportVarConflicts(v *Var) {
	// Ensure that the newly added var does not conflict with a package import
	// which was added earlier.
	if _, ok := m.registry.searchImport(v.Name); ok {
		v.Name += "MoqParam"
	}
	// Ensure that all the newly added imports do not conflict with any of the
	// existing vars.
	for _, imprt := range v.imports {
		if v, ok := m.searchVar(imprt.Qualifier()); ok {
			v.Name += "MoqParam"
		}
	}
}