// Copyright 2023 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 loopvar applies the proper variable capture, according
// to experiment, flags, language version, etc.
package loopvar

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

// ForCapture transforms for and range loops that declare variables that might be
// captured by a closure or escaped to the heap, using a syntactic check that
// conservatively overestimates the loops where capture occurs, but still avoids
// transforming the (large) majority of loops. It returns the list of names
// subject to this change, that may (once transformed) be heap allocated in the
// process. (This allows checking after escape analysis to call out any such
// variables, in case it causes allocation/performance problems).
//
// The decision to transform loops is normally encoded in the For/Range loop node
// field DistinctVars but is also dependent on base.LoopVarHash, and some values
// of base.Debug.LoopVar (which is set per-package).  Decisions encoded in DistinctVars
// are preserved across inlining, so if package a calls b.F and loops in b.F are
// transformed, then they are always transformed, whether b.F is inlined or not.
//
// Per-package, the debug flag settings that affect this transformer:
//
// base.LoopVarHash != nil => use hash setting to govern transformation.
// note that LoopVarHash != nil sets base.Debug.LoopVar to 1 (unless it is >= 11, for testing/debugging).
//
// base.Debug.LoopVar == 11 => transform ALL loops ignoring syntactic/potential escape. Do not log, can be in addition to GOEXPERIMENT.
//
// The effect of GOEXPERIMENT=loopvar is to change the default value (0) of base.Debug.LoopVar to 1 for all packages.
func ForCapture(fn *ir.Func) []*ir.Name {
	// if a loop variable is transformed it is appended to this slice for later logging
	var transformed []*ir.Name

	forCapture := func() {
		seq := 1

		dclFixups := make(map[*ir.Name]ir.Stmt)

		// possibly leaked includes names of declared loop variables that may be leaked;
		// the mapped value is true if the name is *syntactically* leaked, and those loops
		// will be transformed.
		possiblyLeaked := make(map[*ir.Name]bool)

		// these enable an optimization of "escape" under return statements
		loopDepth := 0
		returnInLoopDepth := 0

		// noteMayLeak is called for candidate variables in for range/3-clause, and
		// adds them (mapped to false) to possiblyLeaked.
		noteMayLeak := func(x ir.Node) {
			if n, ok := x.(*ir.Name); ok {
				if n.Type().Kind() == types.TBLANK {
					return
				}
				// default is false (leak candidate, not yet known to leak), but flag can make all variables "leak"
				possiblyLeaked[n] = base.Debug.LoopVar >= 11
			}
		}

		// maybeReplaceVar unshares an iteration variable for a range loop,
		// if that variable was actually (syntactically) leaked,
		// subject to hash-variable debugging.
		maybeReplaceVar := func(k ir.Node, x *ir.RangeStmt) ir.Node {
			if n, ok := k.(*ir.Name); ok && possiblyLeaked[n] {
				if base.LoopVarHash.DebugHashMatchPos(n.Pos()) {
					// Rename the loop key, prefix body with assignment from loop key
					transformed = append(transformed, n)
					tk := typecheck.Temp(n.Type())
					tk.SetTypecheck(1)
					as := ir.NewAssignStmt(x.Pos(), n, tk)
					as.Def = true
					as.SetTypecheck(1)
					x.Body.Prepend(as)
					dclFixups[n] = as
					return tk
				}
			}
			return k
		}

		// scanChildrenThenTransform processes node x to:
		//  1. if x is a for/range w/ DistinctVars, note declared iteration variables possiblyLeaked (PL)
		//  2. search all of x's children for syntactically escaping references to v in PL,
		//     meaning either address-of-v or v-captured-by-a-closure
		//  3. for all v in PL that had a syntactically escaping reference, transform the declaration
		//     and (in case of 3-clause loop) the loop to the unshared loop semantics.
		//  This is all much simpler for range loops; 3-clause loops can have an arbitrary number
		//  of iteration variables and the transformation is more involved, range loops have at most 2.
		var scanChildrenThenTransform func(x ir.Node) bool
		scanChildrenThenTransform = func(n ir.Node) bool {
			switch x := n.(type) {
			case *ir.ClosureExpr:
				if returnInLoopDepth >= loopDepth {
					// This expression is a child of a return, which escapes all loops above
					// the return, but not those between this expression and the return.
					break
				}
				for _, cv := range x.Func.ClosureVars {
					v := cv.Canonical()
					if _, ok := possiblyLeaked[v]; ok {
						possiblyLeaked[v] = true
					}
				}

			case *ir.AddrExpr:
				if returnInLoopDepth >= loopDepth {
					// This expression is a child of a return, which escapes all loops above
					// the return, but not those between this expression and the return.
					break
				}
				// Explicitly note address-taken so that return-statements can be excluded
				y := ir.OuterValue(x.X)
				if y.Op() != ir.ONAME {
					break
				}
				z, ok := y.(*ir.Name)
				if !ok {
					break
				}
				switch z.Class {
				case ir.PAUTO, ir.PPARAM, ir.PPARAMOUT, ir.PAUTOHEAP:
					if _, ok := possiblyLeaked[z]; ok {
						possiblyLeaked[z] = true
					}
				}

			case *ir.ReturnStmt:
				savedRILD := returnInLoopDepth
				returnInLoopDepth = loopDepth
				defer func() { returnInLoopDepth = savedRILD }()

			case *ir.RangeStmt:
				if !(x.Def && x.DistinctVars) {
					// range loop must define its iteration variables AND have distinctVars.
					x.DistinctVars = false
					break
				}
				noteMayLeak(x.Key)
				noteMayLeak(x.Value)
				loopDepth++
				ir.DoChildren(n, scanChildrenThenTransform)
				loopDepth--
				x.Key = maybeReplaceVar(x.Key, x)
				x.Value = maybeReplaceVar(x.Value, x)
				x.DistinctVars = false
				return false

			case *ir.ForStmt:
				if !x.DistinctVars {
					break
				}
				forAllDefInInit(x, noteMayLeak)
				loopDepth++
				ir.DoChildren(n, scanChildrenThenTransform)
				loopDepth--
				var leaked []*ir.Name
				// Collect the leaking variables for the much-more-complex transformation.
				forAllDefInInit(x, func(z ir.Node) {
					if n, ok := z.(*ir.Name); ok && possiblyLeaked[n] {
						// Hash on n.Pos() for most precise failure location.
						if base.LoopVarHash.DebugHashMatchPos(n.Pos()) {
							leaked = append(leaked, n)
						}
					}
				})

				if len(leaked) > 0 {
					// need to transform the for loop just so.

					/* Contrived example, w/ numbered comments from the transformation:
									BEFORE:
										var escape []*int
										for z := 0; z < n; z++ {
											if reason() {
												escape = append(escape, &z)
												continue
											}
											z = z + z
											stuff
										}
									AFTER:
										for z', tmp_first := 0, true; ; { // (4)
											                              // (5) body' follows:
											z := z'                       // (1)
											if tmp_first {tmp_first = false} else {z++} // (6)
											if ! (z < n) { break }        // (7)
											                              // (3, 8) body_continue
											if reason() {
					                            escape = append(escape, &z)
												goto next                 // rewritten continue
											}
											z = z + z
											stuff
										next:                             // (9)
											z' = z                       // (2)
										}

										In the case that the loop contains no increment (z++),
										there is no need for step 6,
										and thus no need to test, update, or declare tmp_first (part of step 4).
										Similarly if the loop contains no exit test (z < n),
										then there is no need for step 7.
					*/

					// Expressed in terms of the input ForStmt
					//
					// 	type ForStmt struct {
					// 	init     Nodes
					// 	Label    *types.Sym
					// 	Cond     Node  // empty if OFORUNTIL
					// 	Post     Node
					// 	Body     Nodes
					// 	HasBreak bool
					// }

					// OFOR: init; loop: if !Cond {break}; Body; Post; goto loop

					// (1) prebody = {z := z' for z in leaked}
					// (2) postbody = {z' = z for z in leaked}
					// (3) body_continue = {body : s/continue/goto next}
					// (4) init' = (init : s/z/z' for z in leaked) + tmp_first := true
					// (5) body' = prebody +        // appears out of order below
					// (6)         if tmp_first {tmp_first = false} else {Post} +
					// (7)         if !cond {break} +
					// (8)         body_continue (3) +
					// (9)         next: postbody (2)
					// (10) cond' = {}
					// (11) post' = {}

					// minor optimizations:
					//   if Post is empty, tmp_first and step 6 can be skipped.
					//   if Cond is empty, that code can also be skipped.

					var preBody, postBody ir.Nodes

					// Given original iteration variable z, what is the corresponding z'
					// that carries the value from iteration to iteration?
					zPrimeForZ := make(map[*ir.Name]*ir.Name)

					// (1,2) initialize preBody and postBody
					for _, z := range leaked {
						transformed = append(transformed, z)

						tz := typecheck.Temp(z.Type())
						tz.SetTypecheck(1)
						zPrimeForZ[z] = tz

						as := ir.NewAssignStmt(x.Pos(), z, tz)
						as.Def = true
						as.SetTypecheck(1)
						preBody.Append(as)
						dclFixups[z] = as

						as = ir.NewAssignStmt(x.Pos(), tz, z)
						as.SetTypecheck(1)
						postBody.Append(as)

					}

					// (3) rewrite continues in body -- rewrite is inplace, so works for top level visit, too.
					label := typecheck.Lookup(fmt.Sprintf(".3clNext_%d", seq))
					seq++
					labelStmt := ir.NewLabelStmt(x.Pos(), label)
					labelStmt.SetTypecheck(1)

					loopLabel := x.Label
					loopDepth := 0
					var editContinues func(x ir.Node) bool
					editContinues = func(x ir.Node) bool {

						switch c := x.(type) {
						case *ir.BranchStmt:
							// If this is a continue targeting the loop currently being rewritten, transform it to an appropriate GOTO
							if c.Op() == ir.OCONTINUE && (loopDepth == 0 && c.Label == nil || loopLabel != nil && c.Label == loopLabel) {
								c.Label = label
								c.SetOp(ir.OGOTO)
							}
						case *ir.RangeStmt, *ir.ForStmt:
							loopDepth++
							ir.DoChildren(x, editContinues)
							loopDepth--
							return false
						}
						ir.DoChildren(x, editContinues)
						return false
					}
					for _, y := range x.Body {
						editContinues(y)
					}
					bodyContinue := x.Body

					// (4) rewrite init
					forAllDefInInitUpdate(x, func(z ir.Node, pz *ir.Node) {
						// note tempFor[n] can be nil if hash searching.
						if n, ok := z.(*ir.Name); ok && possiblyLeaked[n] && zPrimeForZ[n] != nil {
							*pz = zPrimeForZ[n]
						}
					})

					postNotNil := x.Post != nil
					var tmpFirstDcl *ir.AssignStmt
					if postNotNil {
						// body' = prebody +
						// (6)     if tmp_first {tmp_first = false} else {Post} +
						//         if !cond {break} + ...
						tmpFirst := typecheck.Temp(types.Types[types.TBOOL])

						// tmpFirstAssign assigns val to tmpFirst
						tmpFirstAssign := func(val bool) *ir.AssignStmt {
							s := ir.NewAssignStmt(x.Pos(), tmpFirst, typecheck.OrigBool(tmpFirst, val))
							s.SetTypecheck(1)
							return s
						}

						tmpFirstDcl = tmpFirstAssign(true)
						tmpFirstDcl.Def = true // also declares tmpFirst
						tmpFirstSetFalse := tmpFirstAssign(false)
						ifTmpFirst := ir.NewIfStmt(x.Pos(), tmpFirst, ir.Nodes{tmpFirstSetFalse}, ir.Nodes{x.Post})
						ifTmpFirst.SetTypecheck(1)
						preBody.Append(ifTmpFirst)
					}

					// body' = prebody +
					//         if tmp_first {tmp_first = false} else {Post} +
					// (7)     if !cond {break} + ...
					if x.Cond != nil {
						notCond := ir.NewUnaryExpr(x.Cond.Pos(), ir.ONOT, x.Cond)
						notCond.SetType(x.Cond.Type())
						notCond.SetTypecheck(1)
						newBreak := ir.NewBranchStmt(x.Pos(), ir.OBREAK, nil)
						newBreak.SetTypecheck(1)
						ifNotCond := ir.NewIfStmt(x.Pos(), notCond, ir.Nodes{newBreak}, nil)
						ifNotCond.SetTypecheck(1)
						preBody.Append(ifNotCond)
					}

					if postNotNil {
						x.PtrInit().Append(tmpFirstDcl)
					}

					// (8)
					preBody.Append(bodyContinue...)
					// (9)
					preBody.Append(labelStmt)
					preBody.Append(postBody...)

					// (5) body' = prebody + ...
					x.Body = preBody

					// (10) cond' = {}
					x.Cond = nil

					// (11) post' = {}
					x.Post = nil
				}
				x.DistinctVars = false

				return false
			}

			ir.DoChildren(n, scanChildrenThenTransform)

			return false
		}
		scanChildrenThenTransform(fn)
		if len(transformed) > 0 {
			// editNodes scans a slice C of ir.Node, looking for declarations that
			// appear in dclFixups.  Any declaration D whose "fixup" is an assignmnt
			// statement A is removed from the C and relocated to the Init
			// of A.  editNodes returns the modified slice of ir.Node.
			editNodes := func(c ir.Nodes) ir.Nodes {
				j := 0
				for _, n := range c {
					if d, ok := n.(*ir.Decl); ok {
						if s := dclFixups[d.X]; s != nil {
							switch a := s.(type) {
							case *ir.AssignStmt:
								a.PtrInit().Prepend(d)
								delete(dclFixups, d.X) // can't be sure of visit order, wouldn't want to visit twice.
							default:
								base.Fatalf("not implemented yet for node type %v", s.Op())
							}
							continue // do not copy this node, and do not increment j
						}
					}
					c[j] = n
					j++
				}
				for k := j; k < len(c); k++ {
					c[k] = nil
				}
				return c[:j]
			}
			// fixup all tagged declarations in all the statements lists in fn.
			rewriteNodes(fn, editNodes)
		}
	}
	ir.WithFunc(fn, forCapture)
	return transformed
}

// forAllDefInInitUpdate applies "do" to all the defining assignemnts in the Init clause of a ForStmt.
// This abstracts away some of the boilerplate from the already complex and verbose for-3-clause case.
func forAllDefInInitUpdate(x *ir.ForStmt, do func(z ir.Node, update *ir.Node)) {
	for _, s := range x.Init() {
		switch y := s.(type) {
		case *ir.AssignListStmt:
			if !y.Def {
				continue
			}
			for i, z := range y.Lhs {
				do(z, &y.Lhs[i])
			}
		case *ir.AssignStmt:
			if !y.Def {
				continue
			}
			do(y.X, &y.X)
		}
	}
}

// forAllDefInInit is forAllDefInInitUpdate without the update option.
func forAllDefInInit(x *ir.ForStmt, do func(z ir.Node)) {
	forAllDefInInitUpdate(x, func(z ir.Node, _ *ir.Node) { do(z) })
}

// rewriteNodes applies editNodes to all statement lists in fn.
func rewriteNodes(fn *ir.Func, editNodes func(c ir.Nodes) ir.Nodes) {
	var forNodes func(x ir.Node) bool
	forNodes = func(n ir.Node) bool {
		if stmt, ok := n.(ir.InitNode); ok {
			// process init list
			stmt.SetInit(editNodes(stmt.Init()))
		}
		switch x := n.(type) {
		case *ir.Func:
			x.Body = editNodes(x.Body)
			x.Enter = editNodes(x.Enter)
			x.Exit = editNodes(x.Exit)
		case *ir.InlinedCallExpr:
			x.Body = editNodes(x.Body)

		case *ir.CaseClause:
			x.Body = editNodes(x.Body)
		case *ir.CommClause:
			x.Body = editNodes(x.Body)

		case *ir.BlockStmt:
			x.List = editNodes(x.List)

		case *ir.ForStmt:
			x.Body = editNodes(x.Body)
		case *ir.RangeStmt:
			x.Body = editNodes(x.Body)
		case *ir.IfStmt:
			x.Body = editNodes(x.Body)
			x.Else = editNodes(x.Else)
		case *ir.SelectStmt:
			x.Compiled = editNodes(x.Compiled)
		case *ir.SwitchStmt:
			x.Compiled = editNodes(x.Compiled)
		}
		ir.DoChildren(n, forNodes)
		return false
	}
	forNodes(fn)
}