// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements type parameter substitution.

package types2

import (
	"cmd/compile/internal/syntax"
)

type substMap map[*TypeParam]Type

// makeSubstMap creates a new substitution map mapping tpars[i] to targs[i].
// If targs[i] is nil, tpars[i] is not substituted.
func makeSubstMap(tpars []*TypeParam, targs []Type) substMap {
	assert(len(tpars) == len(targs))
	proj := make(substMap, len(tpars))
	for i, tpar := range tpars {
		proj[tpar] = targs[i]
	}
	return proj
}

// makeRenameMap is like makeSubstMap, but creates a map used to rename type
// parameters in from with the type parameters in to.
func makeRenameMap(from, to []*TypeParam) substMap {
	assert(len(from) == len(to))
	proj := make(substMap, len(from))
	for i, tpar := range from {
		proj[tpar] = to[i]
	}
	return proj
}

func (m substMap) empty() bool {
	return len(m) == 0
}

func (m substMap) lookup(tpar *TypeParam) Type {
	if t := m[tpar]; t != nil {
		return t
	}
	return tpar
}

// subst returns the type typ with its type parameters tpars replaced by the
// corresponding type arguments targs, recursively. subst doesn't modify the
// incoming type. If a substitution took place, the result type is different
// from the incoming type.
//
// If expanding is non-nil, it is the instance type currently being expanded.
// One of expanding or ctxt must be non-nil.
func (check *Checker) subst(pos syntax.Pos, typ Type, smap substMap, expanding *Named, ctxt *Context) Type {
	assert(expanding != nil || ctxt != nil)

	if smap.empty() {
		return typ
	}

	// common cases
	switch t := typ.(type) {
	case *Basic:
		return typ // nothing to do
	case *TypeParam:
		return smap.lookup(t)
	}

	// general case
	subst := subster{
		pos:       pos,
		smap:      smap,
		check:     check,
		expanding: expanding,
		ctxt:      ctxt,
	}
	return subst.typ(typ)
}

type subster struct {
	pos       syntax.Pos
	smap      substMap
	check     *Checker // nil if called via Instantiate
	expanding *Named   // if non-nil, the instance that is being expanded
	ctxt      *Context
}

func (subst *subster) typ(typ Type) Type {
	switch t := typ.(type) {
	case nil:
		// Call typOrNil if it's possible that typ is nil.
		panic("nil typ")

	case *Basic:
		// nothing to do

	case *Alias:
		// This code follows the code for *Named types closely.
		// TODO(gri) try to factor better
		orig := t.Origin()
		n := orig.TypeParams().Len()
		if n == 0 {
			return t // type is not parameterized
		}

		// TODO(gri) do we need this for Alias types?
		if t.TypeArgs().Len() != n {
			return Typ[Invalid] // error reported elsewhere
		}

		// already instantiated
		// For each (existing) type argument determine if it needs
		// to be substituted; i.e., if it is or contains a type parameter
		// that has a type argument for it.
		if targs := substList(t.TypeArgs().list(), subst.typ); targs != nil {
			return subst.check.newAliasInstance(subst.pos, t.orig, targs, subst.expanding, subst.ctxt)
		}

	case *Array:
		elem := subst.typOrNil(t.elem)
		if elem != t.elem {
			return &Array{len: t.len, elem: elem}
		}

	case *Slice:
		elem := subst.typOrNil(t.elem)
		if elem != t.elem {
			return &Slice{elem: elem}
		}

	case *Struct:
		if fields := substList(t.fields, subst.var_); fields != nil {
			s := &Struct{fields: fields, tags: t.tags}
			s.markComplete()
			return s
		}

	case *Pointer:
		base := subst.typ(t.base)
		if base != t.base {
			return &Pointer{base: base}
		}

	case *Tuple:
		return subst.tuple(t)

	case *Signature:
		// Preserve the receiver: it is handled during *Interface and *Named type
		// substitution.
		//
		// Naively doing the substitution here can lead to an infinite recursion in
		// the case where the receiver is an interface. For example, consider the
		// following declaration:
		//
		//  type T[A any] struct { f interface{ m() } }
		//
		// In this case, the type of f is an interface that is itself the receiver
		// type of all of its methods. Because we have no type name to break
		// cycles, substituting in the recv results in an infinite loop of
		// recv->interface->recv->interface->...
		recv := t.recv

		// If t is a generic method signature whose own type parameters are not
		// themselves the subject of this substitution, we are substituting the
		// receiver type parameters (via Named.expandMethod). Because a method
		// type parameter's bound may refer to a receiver type parameter
		// (e.g. func (G[T]) M[P interface{ ~*T }]), we must create fresh type
		// parameters with substituted bounds, and rename occurrences in params
		// and results so they refer to the fresh parameters. Otherwise the
		// resulting signature would retain a free reference to the original
		// receiver type parameter.
		//
		// Fresh type parameters are always created, even when the bounds are
		// unaffected by the substitution, so that the methods of distinct
		// instances of the receiver type have distinct (method-specific) type
		// parameters.
		//
		// When t's type parameters are the variables being substituted, we are
		// instantiating t itself; the caller (Checker.instance for *Signature)
		// sets tparams to nil afterward, so we leave them in place here.
		tparams := t.tparams
		s := subst
		if n := tparams.Len(); n > 0 {
			// If (any) one of the signature's type parameters is in the
			// substitution map, this subst call is an instantiation of the
			// signature.
			_, instantiating := subst.smap[tparams.At(0)]
			if debug {
				// When calling subst on a signature, the substitution either
				// applies to all of the type parameters (all are in the map)
				// or none of them (none are in the map).
				for _, tp := range tparams.list() {
					_, ok := subst.smap[tp]
					assert(ok == instantiating)
				}
			}
			if !instantiating {
				fresh := make([]*TypeParam, n)
				// We're introducing a fresh set of method type parameters
				// which appear elsewhere in the signature (parameter or
				// result types, or the bounds of other type parameters).
				// Create an updated substitution map containing the
				// existing entries plus an entry for each fresh type
				// parameter so that they are substituted simultaneously
				// when we proceed with the outer substitution.
				smap := make(substMap, len(subst.smap)+n)
				for k, v := range subst.smap {
					smap[k] = v
				}
				for i, tp := range tparams.list() {
					tname := NewTypeName(tp.Obj().Pos(), tp.Obj().Pkg(), tp.Obj().Name(), nil)
					ftp := subst.check.newTypeParam(tname, nil)
					ftp.index = tp.index
					fresh[i] = ftp
					smap[tp] = ftp
					// The fresh parameter stands in for tp in the mono graph,
					// so that instantiations of (e.g.) G[int].M and G[A].M
					// are tracked against the same vertex as the origin's tp.
					if subst.check != nil {
						subst.check.mono.recordCanon(ftp, tp)
					}
				}
				// Now that we have the updated substitution map, use it to
				// compute the constraints for the fresh type parameters.
				for i, tp := range tparams.list() {
					fresh[i].bound = subst.check.subst(subst.pos, tp.bound, smap, subst.expanding, subst.ctxt)
				}
				// Continue with the fresh type parameters and updated map.
				tparams = &TypeParamList{tparams: fresh}
				s = &subster{
					pos:       subst.pos,
					smap:      smap,
					check:     subst.check,
					expanding: subst.expanding,
					ctxt:      subst.ctxt,
				}
			}
		}

		params := s.tuple(t.params)
		results := s.tuple(t.results)
		if params != t.params || results != t.results || tparams != t.tparams {
			return &Signature{
				rparams: t.rparams,
				tparams: tparams,
				// instantiated signatures have a nil scope
				recv:     recv,
				recvold:  t.recvold,
				params:   params,
				results:  results,
				variadic: t.variadic,
			}
		}

	case *Union:
		if terms := substList(t.terms, subst.term); terms != nil {
			// term list substitution may introduce duplicate terms (unlikely but possible).
			// This is ok; lazy type set computation will determine the actual type set
			// in normal form.
			return &Union{terms}
		}

	case *Interface:
		methods := substList(t.methods, subst.func_)
		embeddeds := substList(t.embeddeds, subst.typ)
		if methods != nil || embeddeds != nil {
			if methods == nil {
				methods = t.methods
			}
			if embeddeds == nil {
				embeddeds = t.embeddeds
			}
			iface := subst.check.newInterface()
			iface.embeddeds = embeddeds
			iface.embedPos = t.embedPos
			iface.implicit = t.implicit
			assert(t.complete) // otherwise we are copying incomplete data
			iface.complete = t.complete
			// If we've changed the interface type, we may need to replace its
			// receiver if the receiver type is the original interface. Receivers of
			// *Named type are replaced during named type expansion.
			//
			// Notably, it's possible to reach here and not create a new *Interface,
			// even though the receiver type may be parameterized. For example:
			//
			//  type T[P any] interface{ m() }
			//
			// In this case the interface will not be substituted here, because its
			// method signatures do not depend on the type parameter P, but we still
			// need to create new interface methods to hold the instantiated
			// receiver. This is handled by Named.expandUnderlying.
			iface.methods, _ = replaceRecvType(methods, t, iface)

			// If check != nil, check.newInterface will have saved the interface for later completion.
			if subst.check == nil { // golang/go#61561: all newly created interfaces must be completed
				iface.typeSet()
			}
			return iface
		}

	case *Map:
		key := subst.typ(t.key)
		elem := subst.typ(t.elem)
		if key != t.key || elem != t.elem {
			return &Map{key: key, elem: elem}
		}

	case *Chan:
		elem := subst.typ(t.elem)
		if elem != t.elem {
			return &Chan{dir: t.dir, elem: elem}
		}

	case *Named:
		// subst is called during expansion, so in this function we need to be
		// careful not to call any methods that would cause t to be expanded: doing
		// so would result in deadlock.
		//
		// So we call t.Origin().TypeParams() rather than t.TypeParams().
		orig := t.Origin()
		n := orig.TypeParams().Len()
		if n == 0 {
			return t // type is not parameterized
		}

		if t.TypeArgs().Len() != n {
			return Typ[Invalid] // error reported elsewhere
		}

		// already instantiated
		// For each (existing) type argument determine if it needs
		// to be substituted; i.e., if it is or contains a type parameter
		// that has a type argument for it.
		if targs := substList(t.TypeArgs().list(), subst.typ); targs != nil {
			// Create a new instance and populate the context to avoid endless
			// recursion. The position used here is irrelevant because validation only
			// occurs on t (we don't call validType on named), but we use subst.pos to
			// help with debugging.
			return subst.check.instance(subst.pos, orig, targs, subst.expanding, subst.ctxt)
		}

	case *TypeParam:
		return subst.smap.lookup(t)

	default:
		panic("unreachable")
	}

	return typ
}

// typOrNil is like typ but if the argument is nil it is replaced with Typ[Invalid].
// A nil type may appear in pathological cases such as type T[P any] []func(_ T([]_))
// where an array/slice element is accessed before it is set up.
func (subst *subster) typOrNil(typ Type) Type {
	if typ == nil {
		return Typ[Invalid]
	}
	return subst.typ(typ)
}

func (subst *subster) var_(v *Var) *Var {
	if v != nil {
		if typ := subst.typ(v.typ); typ != v.typ {
			return cloneVar(v, typ)
		}
	}
	return v
}

func cloneVar(v *Var, typ Type) *Var {
	copy := *v
	copy.typ = typ
	copy.origin = v.Origin()
	return &copy
}

func (subst *subster) tuple(t *Tuple) *Tuple {
	if t != nil {
		if vars := substList(t.vars, subst.var_); vars != nil {
			return &Tuple{vars: vars}
		}
	}
	return t
}

// substList applies subst to each element of the incoming slice.
// If at least one element changes, the result is a new slice with
// all the (possibly updated) elements of the incoming slice;
// otherwise the result it nil. The incoming slice is unchanged.
func substList[T comparable](in []T, subst func(T) T) (out []T) {
	for i, t := range in {
		if u := subst(t); u != t {
			if out == nil {
				// lazily allocate a new slice on first substitution
				out = make([]T, len(in))
				copy(out, in)
			}
			out[i] = u
		}
	}
	return
}

func (subst *subster) func_(f *Func) *Func {
	if f != nil {
		if typ := subst.typ(f.typ); typ != f.typ {
			return cloneFunc(f, typ)
		}
	}
	return f
}

func cloneFunc(f *Func, typ Type) *Func {
	copy := *f
	copy.typ = typ
	copy.origin = f.Origin()
	return &copy
}

func (subst *subster) term(t *Term) *Term {
	if typ := subst.typ(t.typ); typ != t.typ {
		return NewTerm(t.tilde, typ)
	}
	return t
}

// replaceRecvType updates any function receivers that have type old to have
// type new. It does not modify the input slice; if modifications are required,
// the input slice and any affected signatures will be copied before mutating.
//
// The resulting out slice contains the updated functions, and copied reports
// if anything was modified.
func replaceRecvType(in []*Func, old, new Type) (out []*Func, copied bool) {
	out = in
	for i, method := range in {
		sig := method.Signature()
		if sig.recv != nil && sig.recv.Type() == old {
			if !copied {
				// Allocate a new methods slice before mutating for the first time.
				// This is defensive, as we may share methods across instantiations of
				// a given interface type if they do not get substituted.
				out = make([]*Func, len(in))
				copy(out, in)
				copied = true
			}
			newsig := *sig
			newsig.recv = cloneVar(sig.recv, new)
			out[i] = cloneFunc(method, &newsig)
		}
	}
	return
}
