package main

import (
	. "git.smesh.lol/moxie/pkg/types"
)

// checkUseAfterFree: no-op in arena model.
func checkUseAfterFree(f *SSAFunction) {}

// checkParamMutation rejects stores through function parameters.
// Parameters are borrowed references into the caller's arena. Writing
// through them places callee-arena data in caller-owned storage,
// which dangles when the callee's arena dies. Return values are the
// only way to pass data out.
func checkParamMutation(f *SSAFunction) {
	if f == nil || f.Blocks == nil {
		return
	}
	// Skip runtime - it manages arenas directly.
	if f.Pkg != nil && f.Pkg.Pkg != nil {
		path := f.Pkg.Pkg.Path
		if path == "runtime" || path == "unsafe" {
			return
		}
	}
	// Pass 1: find allocs that hold parameter values.
	// The SSA builder stores each param into an alloc:
	//   %alloc = alloca T
	//   store T %param, ptr %alloc
	paramAllocs := map[SSAValue]bool{}
	for _, p := range f.Params {
		paramAllocs[SSAValue(p)] = true
	}
	paramHolders := map[SSAValue]bool{}
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			s, ok := instr.(*SSAStore)
			if !ok {
				continue
			}
			if paramAllocs[s.Val] {
				paramHolders[s.Addr] = true
			}
		}
	}
	// Pass 2: check all stores.
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			s, ok := instr.(*SSAStore)
			if !ok {
				continue
			}
			if paramAllocs[s.Val] {
				continue
			}
			if isParamDerivedEx(s.Addr, paramHolders, 0) {
				name := paramRootName(s.Addr, paramHolders, 0)
				push(cctx.compileErrors, f.name | ": store through borrowed parameter '" | name | "'")
			}
		}
	}
}

func isParamDerivedEx(v SSAValue, holders map[SSAValue]bool, depth int32) (ok bool) {
	if v == nil || depth > 20 {
		return false
	}
	switch x := v.(type) {
	case *SSAParameter:
		return true
	case *SSAFieldAddr:
		return isParamDerivedEx(x.X, holders, depth+1)
	case *SSAIndexAddr:
		return isParamDerivedEx(x.X, holders, depth+1)
	case *SSAUnOp:
		if x.Op == OpMul {
			if holders[x.X] {
				return true
			}
			return isParamDerivedEx(x.X, holders, depth+1)
		}
	case *SSAChangeType:
		return isParamDerivedEx(x.X, holders, depth+1)
	case *SSAConvert:
		return isParamDerivedEx(x.X, holders, depth+1)
	case *SSAPhi:
		for _, edge := range x.Edges {
			if isParamDerivedEx(edge, holders, depth+1) {
				return true
			}
		}
	}
	return false
}

func paramRootName(v SSAValue, holders map[SSAValue]bool, depth int32) (name string) {
	if v == nil || depth > 20 {
		return "?"
	}
	switch x := v.(type) {
	case *SSAParameter:
		return x.SSAName()
	case *SSAFieldAddr:
		return paramRootName(x.X, holders, depth+1)
	case *SSAIndexAddr:
		return paramRootName(x.X, holders, depth+1)
	case *SSAUnOp:
		if x.Op == OpMul {
			return paramRootName(x.X, holders, depth+1)
		}
	case *SSAChangeType:
		return paramRootName(x.X, holders, depth+1)
	case *SSAConvert:
		return paramRootName(x.X, holders, depth+1)
	case *SSAPhi:
		for _, edge := range x.Edges {
			n := paramRootName(edge, holders, depth+1)
			if n != "?" {
				return n
			}
		}
	}
	return "?"
}

// markReturnEscapes promotes stack allocas to heap when their address can
// reach a return value.
func markReturnEscapes(f *SSAFunction) {
	derived := map[SSAValue][]*SSAAlloc{}
	stored := map[*SSAAlloc][]*SSAAlloc{}
	mergeInto := func(dst []*SSAAlloc, srcs []*SSAAlloc) ([]*SSAAlloc, bool) {
		didMerge := false
		for _, s := range srcs {
			if s == nil {
				continue
			}
			found := false
			for _, c := range dst {
				if c == s {
					found = true
					break
				}
			}
			if !found {
				push(dst, s)
				didMerge = true
			}
		}
		return dst, didMerge
	}
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			if a, ok := instr.(*SSAAlloc); ok {
				derived[SSAValue(a)] = []*SSAAlloc{a}
			}
		}
	}
	changed := true
	for changed {
		changed = false
		for _, b := range f.Blocks {
			for _, instr := range b.Instrs {
				var dst SSAValue
				var srcs []*SSAAlloc
				switch i := instr.(type) {
				case *SSASlice:
					dst = i
					srcs = derived[i.X]
				case *SSAFieldAddr:
					dst = i
					srcs = derived[i.X]
				case *SSAIndexAddr:
					dst = i
					srcs = derived[i.X]
				case *SSAChangeType:
					dst = i
					srcs = derived[i.X]
				case *SSAConvert:
					dst = i
					srcs = derived[i.X]
				case *SSAMakeInterface:
					dst = i
					srcs = derived[i.X]
				case *SSAPhi:
					dst = i
					for _, edge := range i.Edges {
						srcs = srcs | derived[edge]
					}
				case *SSAUnOp:
					if i.Op == OpMul {
						dst = i
						for _, a := range derived[i.X] {
							srcs = srcs | stored[a]
						}
					}
				case *SSAStore:
					if d := derived[i.Val]; len(d) > 0 {
						for _, a := range derived[i.Addr] {
							ns, ch := mergeInto(stored[a], d)
							if ch {
								stored[a] = ns
								changed = true
							}
						}
					}
				}
				if dst != nil && len(srcs) > 0 {
					nd, ch := mergeInto(derived[dst], srcs)
					if ch {
						derived[dst] = nd
						changed = true
					}
				}
			}
		}
	}
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			r, ok := instr.(*SSAReturn)
			if !ok {
				continue
			}
			for _, res := range r.Results {
				work := derived[res]
				seen := map[*SSAAlloc]bool{}
				for wi := 0; wi < len(work); wi++ {
					a := work[wi]
					if a == nil {
						continue
					}
					if seen[a] {
						continue
					}
					seen[a] = true
					a.Heap = true
					work = work | stored[a]
				}
			}
		}
	}
}

func allocEscapesViaCall(a *SSAAlloc, f *SSAFunction) (ok bool) {
	escaped := map[SSAValue]bool{}
	escaped[SSAValue(a)] = true
	changed := true
	for changed {
		changed = false
		for _, b := range f.Blocks {
			for _, instr := range b.Instrs {
				switch i := instr.(type) {
				case *SSAFieldAddr:
					if escaped[i.X] {
						if v, ok2 := instr.(SSAValue); ok2 && !escaped[v] {
							escaped[v] = true
							changed = true
						}
					}
				case *SSAMakeInterface:
					if escaped[i.X] {
						if v, ok2 := instr.(SSAValue); ok2 && !escaped[v] {
							escaped[v] = true
							changed = true
						}
					}
				case *SSAStore:
					if escaped[i.Val] || escaped[i.Addr] {
						if escaped[i.Val] && !escaped[i.Addr] {
							escaped[i.Addr] = true
							changed = true
						}
					}
				case *SSAPhi:
					for _, edge := range i.Edges {
						if escaped[edge] {
							if v, ok2 := instr.(SSAValue); ok2 && !escaped[v] {
								escaped[v] = true
								changed = true
							}
							break
						}
					}
				}
			}
		}
	}
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			switch i := instr.(type) {
			case *SSACall:
				for _, arg := range i.Call.Args {
					if escaped[arg] {
						return true
					}
				}
			case *SSAMapUpdate:
				if escaped[i.Key] || escaped[i.Value] || escaped[i.Map] {
					return true
				}
			case *SSASend:
				if escaped[i.X] {
					return true
				}
			case *SSAInvoke:
				if escaped[i.X] {
					return true
				}
				for _, arg := range i.Args {
					if escaped[arg] {
						return true
					}
				}
			}
		}
	}
	return false
}

// Scope dealloc functions: all no-ops in arena model.
func (e *irEmitter) scopeTrackAlloc(ptrReg string) {}
func (e *irEmitter) scopeBeforeReturn(r *SSAReturn) {}
func (e *irEmitter) emitScopeFreesAt(scopeID int32) {}
func (e *irEmitter) emitScopeRelocateOnStore(s *SSAStore, val, valType string) {}