// Copyright 2022 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.

package ssa

import (
	"fmt"
	"go/ast"
	"go/types"

	"golang.org/x/tools/internal/typeparams"
)

// Instances returns all of the instances generated by runtime types for this function in an unspecified order.
//
// Thread-safe.
//
// This is an experimental interface! It may change without warning.
func (prog *Program) _Instances(fn *Function) []*Function {
	if len(fn._TypeParams) == 0 {
		return nil
	}

	prog.methodsMu.Lock()
	defer prog.methodsMu.Unlock()
	return prog.instances[fn].list()
}

// A set of instantiations of a generic function fn.
type instanceSet struct {
	fn        *Function               // len(fn._TypeParams) > 0 and len(fn._TypeArgs) == 0.
	instances map[*typeList]*Function // canonical type arguments to an instance.
	syntax    *ast.FuncDecl           // fn.syntax copy for instantiating after fn is done. nil on synthetic packages.
	info      *types.Info             // fn.pkg.info copy for building after fn is done.. nil on synthetic packages.

	// TODO(taking): Consider ways to allow for clearing syntax and info when done building.
	// May require a public API change as MethodValue can request these be built after prog.Build() is done.
}

func (insts *instanceSet) list() []*Function {
	if insts == nil {
		return nil
	}

	fns := make([]*Function, 0, len(insts.instances))
	for _, fn := range insts.instances {
		fns = append(fns, fn)
	}
	return fns
}

// createInstanceSet adds a new instanceSet for a generic function fn if one does not exist.
//
// Precondition: fn is a package level declaration (function or method).
//
// EXCLUSIVE_LOCKS_ACQUIRED(prog.methodMu)
func (prog *Program) createInstanceSet(fn *Function) {
	assert(len(fn._TypeParams) > 0 && len(fn._TypeArgs) == 0, "Can only create instance sets for generic functions")

	prog.methodsMu.Lock()
	defer prog.methodsMu.Unlock()

	syntax, _ := fn.syntax.(*ast.FuncDecl)
	assert((syntax == nil) == (fn.syntax == nil), "fn.syntax is either nil or a *ast.FuncDecl")

	if _, ok := prog.instances[fn]; !ok {
		prog.instances[fn] = &instanceSet{
			fn:     fn,
			syntax: syntax,
			info:   fn.info,
		}
	}
}

// needsInstance returns an Function that that is the instantiation of fn with the type arguments targs.
//
// Any CREATEd instance is added to cr.
//
// EXCLUSIVE_LOCKS_ACQUIRED(prog.methodMu)
func (prog *Program) needsInstance(fn *Function, targs []types.Type, cr *creator) *Function {
	prog.methodsMu.Lock()
	defer prog.methodsMu.Unlock()

	return prog.instances[fn].lookupOrCreate(targs, cr)
}

// lookupOrCreate returns the instantiation of insts.fn using targs.
// If the instantiation is reported, this is added to cr.
func (insts *instanceSet) lookupOrCreate(targs []types.Type, cr *creator) *Function {
	if insts.instances == nil {
		insts.instances = make(map[*typeList]*Function)
	}

	// canonicalize on a tuple of targs. Sig is not unique.
	//
	// func A[T any]() {
	//   var x T
	//   fmt.Println("%T", x)
	// }
	key := insts.fn.Prog.canon.List(targs)
	if inst, ok := insts.instances[key]; ok {
		return inst
	}
	instance := createInstance(insts.fn, targs, insts.info, insts.syntax, cr)

	// TODO(taking): Allow for the function to be built after monomorphization is supported.
	instance.syntax = nil // treat instance as an external function to prevent building.

	insts.instances[key] = instance
	return instance
}

// createInstance returns an CREATEd instantiation of fn using targs.
//
// Function is added to cr.
func createInstance(fn *Function, targs []types.Type, info *types.Info, syntax ast.Node, cr *creator) *Function {
	prog := fn.Prog
	var sig *types.Signature
	var obj *types.Func
	if recv := fn.Signature.Recv(); recv != nil {
		// method
		// instantiates m with targs and returns a canonical representative for this method.
		m := fn.object.(*types.Func)
		recv := recvType(m)
		if p, ok := recv.(*types.Pointer); ok {
			recv = p.Elem()
		}
		named := recv.(*types.Named)
		inst, err := typeparams.Instantiate(prog.ctxt, typeparams.NamedTypeOrigin(named), targs, false)
		if err != nil {
			panic(err)
		}
		canon, _, _ := types.LookupFieldOrMethod(prog.canon.Type(inst), true, m.Pkg(), m.Name())
		obj = canon.(*types.Func)
		sig = obj.Type().(*types.Signature)
	} else {
		instSig, err := typeparams.Instantiate(prog.ctxt, fn.Signature, targs, false)
		if err != nil {
			panic(err)
		}
		instance, ok := instSig.(*types.Signature)
		if !ok {
			panic("Instantiate of a Signature returned a non-signature")
		}
		obj = fn.object.(*types.Func) // instantiation does not exist yet
		sig = prog.canon.Type(instance).(*types.Signature)
	}

	name := fmt.Sprintf("%s[%s]", fn.Name(), targs) // may not be unique
	synthetic := fmt.Sprintf("instantiation of %s", fn.Name())
	instance := &Function{
		name:        name,
		object:      obj,
		Signature:   sig,
		Synthetic:   synthetic,
		syntax:      syntax, // on synthetic packages syntax is nil.
		_Origin:     fn,
		pos:         obj.Pos(),
		Pkg:         nil,
		Prog:        fn.Prog,
		_TypeParams: fn._TypeParams,
		_TypeArgs:   targs,
		info:        info, // on synthetic packages info is nil.
		subst:       makeSubster(prog.ctxt, fn._TypeParams, targs, false),
	}
	cr.Add(instance)
	return instance
}