// Copyright 2025 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 bloop

// This file contains support routines for keeping
// statements alive
// in such loops (example):
//
//	for b.Loop() {
//		var a, b int
//		a = 5
//		b = 6
//		f(a, b)
//	}
//
// The results of a, b and f(a, b) will be kept alive.
//
// Formally, the lhs (if they are [ir.Name]-s) of
// [ir.AssignStmt], [ir.AssignListStmt],
// [ir.AssignOpStmt], and the results of [ir.CallExpr]
// or its args if it doesn't return a value will be kept
// alive.
//
// The keep alive logic is implemented with as wrapping a
// runtime.KeepAlive around the Name.
//
// TODO: currently this is implemented with KeepAlive
// because it will prevent DSE and DCE which is probably
// what we want right now. And KeepAlive takes an ssa
// value instead of a symbol, which is easier to manage.
// But since KeepAlive's context was mainly in the runtime
// and GC, should we implement a new intrinsic that lowers
// to OpVarLive? Peeling out the symbols is a bit tricky
// and also VarLive seems to assume that there exists a
// VarDef on the same symbol that dominates it.

import (
	"cmd/compile/internal/base"
	"cmd/compile/internal/ir"
	"cmd/compile/internal/reflectdata"
	"cmd/compile/internal/typecheck"
	"cmd/compile/internal/types"
	"fmt"
)

// getNameFromNode tries to iteratively peel down the node to
// get the name.
func getNameFromNode(n ir.Node) *ir.Name {
	var ret *ir.Name
	if n.Op() == ir.ONAME {
		ret = n.(*ir.Name)
	} else {
		// avoid infinite recursion on circular referencing nodes.
		seen := map[ir.Node]bool{n: true}
		var findName func(ir.Node) bool
		findName = func(a ir.Node) bool {
			if a.Op() == ir.ONAME {
				ret = a.(*ir.Name)
				return true
			}
			if !seen[a] {
				seen[a] = true
				return ir.DoChildren(a, findName)
			}
			return false
		}
		ir.DoChildren(n, findName)
	}
	return ret
}

// keepAliveAt returns a statement that is either curNode, or a
// block containing curNode followed by a call to runtime.keepAlive for each
// ONAME in ns. These calls ensure that names in ns will be live until
// after curNode's execution.
func keepAliveAt(ns []*ir.Name, curNode ir.Node) ir.Node {
	if len(ns) == 0 {
		return curNode
	}

	pos := curNode.Pos()
	calls := []ir.Node{curNode}
	for _, n := range ns {
		if n == nil {
			continue
		}
		if n.Sym() == nil {
			continue
		}
		if n.Sym().IsBlank() {
			continue
		}
		arg := ir.NewConvExpr(pos, ir.OCONV, types.Types[types.TINTER], n)
		if !n.Type().IsInterface() {
			srcRType0 := reflectdata.TypePtrAt(pos, n.Type())
			arg.TypeWord = srcRType0
			arg.SrcRType = srcRType0
		}
		callExpr := typecheck.Call(pos,
			typecheck.LookupRuntime("KeepAlive"),
			[]ir.Node{arg}, false).(*ir.CallExpr)
		callExpr.IsCompilerVarLive = true
		callExpr.NoInline = true
		calls = append(calls, callExpr)
	}

	return ir.NewBlockStmt(pos, calls)
}

func debugName(name *ir.Name, line string) {
	if base.Flag.LowerM > 0 {
		if name.Linksym() != nil {
			fmt.Printf("%v: %s will be kept alive\n", line, name.Linksym().Name)
		} else {
			fmt.Printf("%v: expr will be kept alive\n", line)
		}
	}
}

// preserveStmt transforms stmt so that any names defined/assigned within it
// are used after stmt's execution, preventing their dead code elimination
// and dead store elimination. The return value is the transformed statement.
func preserveStmt(curFn *ir.Func, stmt ir.Node) (ret ir.Node) {
	ret = stmt
	switch n := stmt.(type) {
	case *ir.AssignStmt:
		// Peel down struct and slice indexing to get the names
		name := getNameFromNode(n.X)
		if name != nil {
			debugName(name, ir.Line(stmt))
			ret = keepAliveAt([]*ir.Name{name}, n)
		}
	case *ir.AssignListStmt:
		names := []*ir.Name{}
		for _, lhs := range n.Lhs {
			name := getNameFromNode(lhs)
			if name != nil {
				debugName(name, ir.Line(stmt))
				names = append(names, name)
			}
		}
		ret = keepAliveAt(names, n)
	case *ir.AssignOpStmt:
		name := getNameFromNode(n.X)
		if name != nil {
			debugName(name, ir.Line(stmt))
			ret = keepAliveAt([]*ir.Name{name}, n)
		}
	case *ir.CallExpr:
		names := []*ir.Name{}
		curNode := stmt
		if n.Fun != nil && n.Fun.Type() != nil && n.Fun.Type().NumResults() != 0 {
			// This function's results are not assigned, assign them to
			// auto tmps and then keepAliveAt these autos.
			// Note: markStmt assumes the context that it's called - this CallExpr is
			// not within another OAS2, which is guaranteed by the case above.
			results := n.Fun.Type().Results()
			lhs := make([]ir.Node, len(results))
			for i, res := range results {
				tmp := typecheck.TempAt(n.Pos(), curFn, res.Type)
				lhs[i] = tmp
				names = append(names, tmp)
			}

			// Create an assignment statement.
			assign := typecheck.AssignExpr(
				ir.NewAssignListStmt(n.Pos(), ir.OAS2, lhs,
					[]ir.Node{n})).(*ir.AssignListStmt)
			assign.Def = true
			curNode = assign
			plural := ""
			if len(results) > 1 {
				plural = "s"
			}
			if base.Flag.LowerM > 0 {
				fmt.Printf("%v: function result%s will be kept alive\n", ir.Line(stmt), plural)
			}
		} else {
			// This function probably doesn't return anything, keep its args alive.
			argTmps := []ir.Node{}
			for i, a := range n.Args {
				if name := getNameFromNode(a); name != nil {
					// If they are name, keep them alive directly.
					debugName(name, ir.Line(stmt))
					names = append(names, name)
				} else if a.Op() == ir.OSLICELIT {
					// variadic args are encoded as slice literal.
					s := a.(*ir.CompLitExpr)
					ns := []*ir.Name{}
					for i, n := range s.List {
						if name := getNameFromNode(n); name != nil {
							debugName(name, ir.Line(a))
							ns = append(ns, name)
						} else {
							// We need a temporary to save this arg.
							tmp := typecheck.TempAt(n.Pos(), curFn, n.Type())
							argTmps = append(argTmps, typecheck.AssignExpr(ir.NewAssignStmt(n.Pos(), tmp, n)))
							names = append(names, tmp)
							s.List[i] = tmp
							if base.Flag.LowerM > 0 {
								fmt.Printf("%v: function arg will be kept alive\n", ir.Line(n))
							}
						}
					}
					names = append(names, ns...)
				} else {
					// expressions, we need to assign them to temps and change the original arg to reference
					// them.
					tmp := typecheck.TempAt(n.Pos(), curFn, a.Type())
					argTmps = append(argTmps, typecheck.AssignExpr(ir.NewAssignStmt(n.Pos(), tmp, a)))
					names = append(names, tmp)
					n.Args[i] = tmp
					if base.Flag.LowerM > 0 {
						fmt.Printf("%v: function arg will be kept alive\n", ir.Line(stmt))
					}
				}
			}
			if len(argTmps) > 0 {
				argTmps = append(argTmps, n)
				curNode = ir.NewBlockStmt(n.Pos(), argTmps)
			}
		}
		ret = keepAliveAt(names, curNode)
	}
	return
}

func preserveStmts(curFn *ir.Func, list ir.Nodes) {
	for i := range list {
		list[i] = preserveStmt(curFn, list[i])
	}
}

// isTestingBLoop returns true if it matches the node as a
// testing.(*B).Loop. See issue #61515.
func isTestingBLoop(t ir.Node) bool {
	if t.Op() != ir.OFOR {
		return false
	}
	nFor, ok := t.(*ir.ForStmt)
	if !ok || nFor.Cond == nil || nFor.Cond.Op() != ir.OCALLFUNC {
		return false
	}
	n, ok := nFor.Cond.(*ir.CallExpr)
	if !ok || n.Fun == nil || n.Fun.Op() != ir.OMETHEXPR {
		return false
	}
	name := ir.MethodExprName(n.Fun)
	if name == nil {
		return false
	}
	if fSym := name.Sym(); fSym != nil && name.Class == ir.PFUNC && fSym.Pkg != nil &&
		fSym.Name == "(*B).Loop" && fSym.Pkg.Path == "testing" {
		// Attempting to match a function call to testing.(*B).Loop
		return true
	}
	return false
}

type editor struct {
	inBloop bool
	curFn   *ir.Func
}

func (e editor) edit(n ir.Node) ir.Node {
	e.inBloop = isTestingBLoop(n) || e.inBloop
	// It's in bloop, mark the stmts with bodies.
	ir.EditChildren(n, e.edit)
	if e.inBloop {
		switch n := n.(type) {
		case *ir.ForStmt:
			preserveStmts(e.curFn, n.Body)
		case *ir.IfStmt:
			preserveStmts(e.curFn, n.Body)
			preserveStmts(e.curFn, n.Else)
		case *ir.BlockStmt:
			preserveStmts(e.curFn, n.List)
		case *ir.CaseClause:
			preserveStmts(e.curFn, n.List)
			preserveStmts(e.curFn, n.Body)
		case *ir.CommClause:
			preserveStmts(e.curFn, n.Body)
		}
	}
	return n
}

// BloopWalk performs a walk on all functions in the package
// if it imports testing and wrap the results of all qualified
// statements in a runtime.KeepAlive intrinsic call. See package
// doc for more details.
//
//	for b.Loop() {...}
//
// loop's body.
func BloopWalk(pkg *ir.Package) {
	hasTesting := false
	for _, i := range pkg.Imports {
		if i.Path == "testing" {
			hasTesting = true
			break
		}
	}
	if !hasTesting {
		return
	}
	for _, fn := range pkg.Funcs {
		e := editor{false, fn}
		ir.EditChildren(fn, e.edit)
	}
}