package emit

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

func (e *irEmitter) instrOperands(instr ssa.SSAInstruction) []ssa.SSAValue {
	switch i := instr.(type) {
	case *ssa.SSAStore:
		return []ssa.SSAValue{i.Addr, i.Val}
	case *ssa.SSAUnOp:
		return []ssa.SSAValue{i.X}
	case *ssa.SSABinOp:
		return []ssa.SSAValue{i.X, i.Y}
	case *ssa.SSACall:
		out := []ssa.SSAValue{i.Call.Value}
		for _, a := range i.Call.Args {
			out = append(out, a)
		}
		return out
	case *ssa.SSAFieldAddr:
		return []ssa.SSAValue{i.X}
	case *ssa.SSAIndexAddr:
		return []ssa.SSAValue{i.X, i.Index}
	case *ssa.SSAExtract:
		return []ssa.SSAValue{i.Tuple}
	case *ssa.SSAPhi:
		return i.Edges
	case *ssa.SSAReturn:
		var out []ssa.SSAValue
		for _, r := range i.Results {
			out = append(out, r)
		}
		return out
	case *ssa.SSAIf:
		return []ssa.SSAValue{i.Cond}
	case *ssa.SSAConvert:
		return []ssa.SSAValue{i.X}
	case *ssa.SSAChangeType:
		return []ssa.SSAValue{i.X}
	case *ssa.SSAMakeInterface:
		return []ssa.SSAValue{i.X}
	case *ssa.SSATypeAssert:
		return []ssa.SSAValue{i.X}
	case *ssa.SSASlice:
		out := []ssa.SSAValue{i.X}
		if i.Low != nil { out = append(out, i.Low) }
		if i.High != nil { out = append(out, i.High) }
		if i.Max != nil { out = append(out, i.Max) }
		return out
	case *ssa.SSAMapUpdate:
		return []ssa.SSAValue{i.Map, i.Key, i.Value}
	case *ssa.SSALookup:
		return []ssa.SSAValue{i.X, i.Index}
	case *ssa.SSARange:
		return []ssa.SSAValue{i.X}
	case *ssa.SSANext:
		return []ssa.SSAValue{i.Iter}
	case *ssa.SSASend:
		return []ssa.SSAValue{i.Chan, i.X}
	case *ssa.SSAMakeSlice:
		out := []ssa.SSAValue{i.Len}
		if i.Cap != nil { out = append(out, i.Cap) }
		if i.Data != nil { out = append(out, i.Data) }
		return out
	}
	return nil
}

func simpleItoa(n int32) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	buf := []byte{:0:20}
	for n > 0 {
		buf = append(buf, byte('0'+n%10))
		n /= 10
	}
	if neg {
		buf = append(buf, '-')
	}
	for i, j := 0, len(buf)-1; i < j; i, j = i+1, j-1 {
		buf[i], buf[j] = buf[j], buf[i]
	}
	return string(buf)
}

func (e *irEmitter) deallocSetup(f *ssa.SSAFunction) {
	heapCount := int32(0)
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			if a, ok := instr.(*ssa.SSAAlloc); ok && a.Heap {
				heapCount++
			}
		}
	}
	if heapCount == 0 {
		return
	}
	e.deallocActive = true
	e.deallocAllocs = map[ssa.SSAValue]int32{}
	e.w("  %dealloc.head = alloca ptr\n")
	e.w("  store ptr null, ptr %dealloc.head\n")
	idx := int32(0)
	for _, b := range f.Blocks {
		for _, instr := range b.Instrs {
			if a, ok := instr.(*ssa.SSAAlloc); ok && a.Heap {
				e.deallocAllocs[a] = idx
				e.w("  %dealloc.entry." | irItoa(idx) | " = alloca {ptr, ptr}\n")
				e.w("  store {ptr, ptr} zeroinitializer, ptr %dealloc.entry." | irItoa(idx) | "\n")
				idx++
			}
		}
	}
	e.deallocCount = idx
	e.declareRuntime("runtime.deallocPush", "void", "ptr, ptr, ptr")
	e.declareRuntime("runtime.deallocFree", "void", "ptr")
	e.declareRuntime("runtime.deallocEscape", "void", "ptr, ptr")
}

func (e *irEmitter) deallocAfterAlloc(a *ssa.SSAAlloc, reg string) {
	if !e.deallocActive {
		return
	}
	idx, ok := e.deallocAllocs[a]
	if !ok {
		return
	}
	entry := "%dealloc.entry." | irItoa(idx)
	e.w("  call void @runtime.deallocPush(ptr " | reg | ", ptr %dealloc.head, ptr " | entry | ")\n")
}

func (e *irEmitter) deallocBeforeReturn(r *ssa.SSAReturn) {
	if !e.deallocActive {
		return
	}
	for _, res := range r.Results {
		e.deallocEscapeResult(res)
	}
	e.nextReg++
	tmp := "%dealloc.list" | irItoa(e.nextReg)
	e.w("  " | tmp | " = load ptr, ptr %dealloc.head\n")
	e.w("  call void @runtime.deallocFree(ptr " | tmp | ")\n")
}

func (e *irEmitter) deallocEscapeResult(v ssa.SSAValue) {
	if v == nil {
		return
	}
	if a, ok := v.(*ssa.SSAAlloc); ok {
		if _, tracked := e.deallocAllocs[a]; tracked {
			e.nextReg++
			tmp := "%dealloc.esc.list" | irItoa(e.nextReg)
			e.w("  " | tmp | " = load ptr, ptr %dealloc.head\n")
			e.w("  call void @runtime.deallocEscape(ptr " | e.operand(v) | ", ptr " | tmp | ")\n")
		}
	}
}

func (e *irEmitter) scopeTrackAlloc(ptrReg string) {
	if !e.scopeActive {
		return
	}
	e.w("  call void @runtime.scopePush(ptr ")
	e.w(ptrReg)
	e.w(", ptr %scope.head)\n")
}

func (e *irEmitter) scopeEscapeReturnResults(r *ssa.SSAReturn) {
	if !e.scopeActive {
		return
	}
	for _, res := range r.Results {
		if res == nil {
			continue
		}
		typ := e.llvmType(res.SSAType())
		val := e.operand(res)
		if val == "null" || val == "zeroinitializer" || val == "0" {
			continue
		}
		ptrReg := ""
		sty := e.sliceType()
		if typ == "ptr" {
			ptrReg = val
		} else if typ == sty {
			e.nextReg++
			ptrReg = "%esc.p" | irItoa(e.nextReg)
			e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue ") ; e.w(sty) ; e.w(" ") ; e.w(val) ; e.w(", 0\n")
		} else if typ == "{ptr, ptr}" {
			e.nextReg++
			ptrReg = "%esc.p" | irItoa(e.nextReg)
			e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue {ptr, ptr} ") ; e.w(val) ; e.w(", 0\n")
		}
		if ptrReg == "" {
			continue
		}
		e.w("  call void @runtime.scopeEscapeReturn(ptr ")
		e.w(ptrReg)
		e.w(", ptr %scope.head)\n")
		e.declareRuntime("runtime.scopeEscapeReturn", "void", "ptr, ptr")
	}
	if len(e.curFunc.NamedResults) > 0 {
		for _, nr := range e.curFunc.NamedResults {
			nrt := e.namedResultType(nr)
			if nrt != "ptr" && nrt != e.sliceType() && nrt != "{ptr, ptr}" {
				continue
			}
			e.nextReg++
			tmp := "%esc.nr" | irItoa(e.nextReg)
			e.w("  ") ; e.w(tmp) ; e.w(" = load ") ; e.w(nrt) ; e.w(", ptr ") ; e.w(e.regName(nr)) ; e.w("\n")
			ptrReg := ""
			sty := e.sliceType()
			if nrt == "ptr" {
				ptrReg = tmp
			} else if nrt == sty {
				e.nextReg++
				ptrReg = "%esc.nrp" | irItoa(e.nextReg)
				e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue ") ; e.w(sty) ; e.w(" ") ; e.w(tmp) ; e.w(", 0\n")
			} else if nrt == "{ptr, ptr}" {
				e.nextReg++
				ptrReg = "%esc.nrp" | irItoa(e.nextReg)
				e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue {ptr, ptr} ") ; e.w(tmp) ; e.w(", 0\n")
			}
			if ptrReg == "" {
				continue
			}
			e.w("  call void @runtime.scopeEscapeReturn(ptr ")
			e.w(ptrReg)
			e.w(", ptr %scope.head)\n")
			e.declareRuntime("runtime.scopeEscapeReturn", "void", "ptr, ptr")
		}
	}
}

func (e *irEmitter) emitScopeFreeAll() {
	if !e.scopeActive {
		return
	}
	e.nextReg++
	cur := "%sf.all" | irItoa(e.nextReg)
	e.w("  ")
	e.w(cur)
	e.w(" = load ptr, ptr %scope.head\n")
	e.nextReg++
	fn := "%sf.fn" | irItoa(e.nextReg)
	e.w("  ")
	e.w(fn)
	e.w(" = call ptr @runtime.scopeFreeTo(ptr ")
	e.w(cur)
	e.w(", ptr null)\n")
	e.w("  store ptr ")
	e.w(fn)
	e.w(", ptr %scope.head\n")
}

func (e *irEmitter) emitScopeFreesTo(scopeID int32) {
	e.nextReg++
	saved := "%sf.sv" | irItoa(e.nextReg)
	e.w("  ")
	e.w(saved)
	e.w(" = load ptr, ptr %scope.save.")
	e.w(irItoa(scopeID))
	e.w("\n")
	e.nextReg++
	cur := "%sf.cur" | irItoa(e.nextReg)
	e.w("  ")
	e.w(cur)
	e.w(" = load ptr, ptr %scope.head\n")
	e.nextReg++
	fn := "%sf.fn" | irItoa(e.nextReg)
	e.w("  ")
	e.w(fn)
	e.w(" = call ptr @runtime.scopeFreeTo(ptr ")
	e.w(cur)
	e.w(", ptr ")
	e.w(saved)
	e.w(")\n")
	e.w("  store ptr ")
	e.w(fn)
	e.w(", ptr %scope.head\n")
}

func (e *irEmitter) emitScopeSave(b *ssa.SSABasicBlock) {
	if b.ScopeID <= 0 {
		return
	}
	enteringScope := false
	for _, pred := range b.Preds {
		if pred.ScopeID != b.ScopeID {
			enteringScope = true
			break
		}
	}
	if enteringScope || len(b.Preds) == 0 {
		e.nextReg++
		tmp := "%ss" | irItoa(e.nextReg)
		e.w("  ")
		e.w(tmp)
		e.w(" = load ptr, ptr %scope.head\n")
		e.w("  store ptr ")
		e.w(tmp)
		e.w(", ptr %scope.save.")
		e.w(irItoa(b.ScopeID))
		e.w("\n")
	}
}

func (e *irEmitter) emitScopeRelocateOnStore(s *ssa.SSAStore, val, valType string) {
	if !e.scopeActive {
		return
	}
	blk := s.InstrBlock()
	if blk == nil || blk.ScopeID == 0 {
		return
	}
	addrScope := int32(0)
	if aid, ok := e.instrScope[s.Addr]; ok {
		addrScope = aid
	}
	if addrScope >= blk.ScopeID {
		return
	}
	if val == "null" || val == "zeroinitializer" || val == "0" {
		return
	}
	ptrReg := ""
	sty := e.sliceType()
	if valType == "ptr" {
		ptrReg = val
	} else if valType == sty {
		e.nextReg++
		ptrReg = "%sr.p" | irItoa(e.nextReg)
		e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue ") ; e.w(sty) ; e.w(" ") ; e.w(val) ; e.w(", 0\n")
	} else if valType == "{ptr, ptr}" {
		e.nextReg++
		ptrReg = "%sr.p" | irItoa(e.nextReg)
		e.w("  ") ; e.w(ptrReg) ; e.w(" = extractvalue {ptr, ptr} ") ; e.w(val) ; e.w(", 0\n")
	}
	if ptrReg == "" {
		return
	}
	e.w("  call void @runtime.scopeRelocate(ptr ")
	e.w(ptrReg)
	e.w(", ptr %scope.head, ptr %scope.save.")
	e.w(irItoa(blk.ScopeID))
	e.w(")\n")
	e.declareRuntime("runtime.scopeRelocate", "void", "ptr, ptr, ptr")
}