// Package regalloc performs register allocation. The algorithm can work on any ISA by implementing the interfaces in
// api.go.
//
// References:
//   - https://web.stanford.edu/class/archive/cs/cs143/cs143.1128/lectures/17/Slides17.pdf
//   - https://en.wikipedia.org/wiki/Chaitin%27s_algorithm
//   - https://llvm.org/ProjectsWithLLVM/2004-Fall-CS426-LS.pdf
//   - https://pfalcon.github.io/ssabook/latest/book-full.pdf: Chapter 9. for liveness analysis.
//   - https://github.com/golang/go/blob/release-branch.go1.21/src/cmd/compile/internal/ssa/regalloc.go
package regalloc

import (
	
	
	

	
)

// NewAllocator returns a new Allocator.
func [ Instr,  Block[],  Function[, ]]( *RegisterInfo) Allocator[, , ] {
	 := Allocator[, , ]{
		regInfo:            ,
		phiDefInstListPool: wazevoapi.NewPool[phiDefInstList[]](resetPhiDefInstList[]),
		blockStates:        wazevoapi.NewIDedPool[blockState[, , ]](resetBlockState[, , ]),
	}
	.state.vrStates = wazevoapi.NewIDedPool[vrState[, , ]](resetVrState[, , ])
	.state.reset()
	for ,  := range .AllocatableRegisters {
		for ,  := range  {
			.allocatableSet = .allocatableSet.add()
		}
	}
	return 
}

type (
	// RegisterInfo holds the statically-known ISA-specific register information.
	RegisterInfo struct {
		// AllocatableRegisters is a 2D array of allocatable RealReg, indexed by regTypeNum and regNum.
		// The order matters: the first element is the most preferred one when allocating.
		AllocatableRegisters [NumRegType][]RealReg
		CalleeSavedRegisters RegSet
		CallerSavedRegisters RegSet
		RealRegToVReg        []VReg
		// RealRegName returns the name of the given RealReg for debugging.
		RealRegName func(r RealReg) string
		RealRegType func(r RealReg) RegType
	}

	// Allocator is a register allocator.
	Allocator[ Instr,  Block[],  Function[, ]] struct {
		// regInfo is static per ABI/ISA, and is initialized by the machine during Machine.PrepareRegisterAllocator.
		regInfo *RegisterInfo
		// allocatableSet is a set of allocatable RealReg derived from regInfo. Static per ABI/ISA.
		allocatableSet           RegSet
		allocatedCalleeSavedRegs []VReg
		vs                       []VReg
		ss                       []*vrState[, , ]
		copies                   []_copy[, , ]
		phiDefInstListPool       wazevoapi.Pool[phiDefInstList[]]

		// Followings are re-used during various places.
		blks  []
		reals []RealReg

		// Following two fields are updated while iterating the blocks in the reverse postorder.
		state       state[, , ]
		blockStates wazevoapi.IDedPool[blockState[, , ]]
	}

	// _copy represents a source and destination pair of a copy instruction.
	_copy[ Instr,  Block[],  Function[, ]] struct {
		src   *vrState[, , ]
		dstID VRegID
	}

	// programCounter represents an opaque index into the program which is used to represents a LiveInterval of a VReg.
	programCounter int32

	state[ Instr,  Block[],  Function[, ]] struct {
		argRealRegs []VReg
		regsInUse   regInUseSet[, , ]
		vrStates    wazevoapi.IDedPool[vrState[, , ]]

		currentBlockID int32

		// allocatedRegSet is a set of RealReg that are allocated during the allocation phase. This is reset per function.
		allocatedRegSet RegSet
	}

	blockState[ Instr,  Block[],  Function[, ]] struct {
		// liveIns is a list of VReg that are live at the beginning of the block.
		liveIns []*vrState[, , ]
		// seen is true if the block is visited during the liveness analysis.
		seen bool
		// visited is true if the block is visited during the allocation phase.
		visited            bool
		startFromPredIndex int
		// startRegs is a list of RealReg that are used at the beginning of the block. This is used to fix the merge edges.
		startRegs regInUseSet[, , ]
		// endRegs is a list of RealReg that are used at the end of the block. This is used to fix the merge edges.
		endRegs regInUseSet[, , ]
	}

	vrState[ Instr,  Block[],  Function[, ]] struct {
		v VReg
		r RealReg
		// defInstr is the instruction that defines this value. If this is the phi value and not the entry block, this is nil.
		defInstr 
		// defBlk is the block that defines this value. If this is the phi value, this is the block whose arguments contain this value.
		defBlk 
		// lca = lowest common ancestor. This is the block that is the lowest common ancestor of all the blocks that
		// reloads this value. This is used to determine the spill location. Only valid if spilled=true.
		lca 
		// lastUse is the program counter of the last use of this value. This changes while iterating the block, and
		// should not be used across the blocks as it becomes invalid. To check the validity, use lastUseUpdatedAtBlockID.
		lastUse                 programCounter
		lastUseUpdatedAtBlockID int32
		// spilled is true if this value is spilled i.e. the value is reload from the stack somewhere in the program.
		//
		// Note that this field is used during liveness analysis for different purpose. This is used to determine the
		// value is live-in or not.
		spilled bool
		// isPhi is true if this is a phi value.
		isPhi      bool
		desiredLoc desiredLoc
		// phiDefInstList is a list of instructions that defines this phi value.
		// This is used to determine the spill location, and only valid if isPhi=true.
		*phiDefInstList[]
	}

	// phiDefInstList is a linked list of instructions that defines a phi value.
	phiDefInstList[ Instr] struct {
		instr 
		v     VReg
		next  *phiDefInstList[]
	}

	// desiredLoc represents a desired location for a VReg.
	desiredLoc uint16
	// desiredLocKind is a kind of desired location for a VReg.
	desiredLocKind uint16
)

const (
	// desiredLocKindUnspecified is a kind of desired location for a VReg that is not specified.
	desiredLocKindUnspecified desiredLocKind = iota
	// desiredLocKindStack is a kind of desired location for a VReg that is on the stack, only used for the phi values.
	desiredLocKindStack
	// desiredLocKindReg is a kind of desired location for a VReg that is in a register.
	desiredLocKindReg
	desiredLocUnspecified = desiredLoc(desiredLocKindUnspecified)
	desiredLocStack       = desiredLoc(desiredLocKindStack)
)

func newDesiredLocReg( RealReg) desiredLoc {
	return desiredLoc(desiredLocKindReg) | desiredLoc(<<2)
}

func ( desiredLoc) () RealReg {
	return RealReg( >> 2)
}

func ( desiredLoc) () bool {
	return &3 == desiredLoc(desiredLocKindStack)
}

func resetPhiDefInstList[ Instr]( *phiDefInstList[]) {
	var  
	.instr = 
	.next = nil
	.v = VRegInvalid
}

func ( *state[, , ]) ( *RegisterInfo) { //nolint:unused
	fmt.Println("\t\tstate:")
	fmt.Println("\t\t\targRealRegs:", .argRealRegs)
	fmt.Println("\t\t\tregsInUse", .regsInUse.format())
	fmt.Println("\t\t\tallocatedRegSet:", .allocatedRegSet.format())
	fmt.Println("\t\t\tused:", .regsInUse.format())
	var  []string
	for  := 0;  <= .vrStates.MaxIDEncountered(); ++ {
		 := .vrStates.Get()
		if  == nil {
			continue
		}
		if .r != RealRegInvalid {
			 = append(, fmt.Sprintf("(v%d: %s)", .v.ID(), .RealRegName(.r)))
		}
	}
	fmt.Println("\t\t\tvrStates:", strings.Join(, ", "))
}

func ( *state[, , ]) () {
	.argRealRegs = .argRealRegs[:0]
	.vrStates.Reset()
	.allocatedRegSet = RegSet(0)
	.regsInUse.reset()
	.currentBlockID = -1
}

func resetVrState[ Instr,  Block[],  Function[, ]]( *vrState[, , ]) {
	.v = VRegInvalid
	.r = RealRegInvalid
	var  
	.defInstr = 
	var  
	.defBlk = 
	.spilled = false
	.lastUse = -1
	.lastUseUpdatedAtBlockID = -1
	.lca = 
	.isPhi = false
	.phiDefInstList = nil
	.desiredLoc = desiredLocUnspecified
}

func ( *state[, , ]) ( VReg) *vrState[, , ] {
	 := .vrStates.GetOrAllocate(int(.ID()))
	if .v == VRegInvalid {
		.v = 
	}
	return 
}

func ( *state[, , ]) ( VRegID) *vrState[, , ] {
	return .vrStates.Get(int())
}

func ( *state[, , ]) ( RealReg,  *vrState[, , ]) {
	.regsInUse.add(, )
	.r = 
	.allocatedRegSet = .allocatedRegSet.add()
}

func ( *state[, , ]) ( RealReg) {
	 := .regsInUse.get()
	if  != nil {
		.regsInUse.remove()
		.r = RealRegInvalid
	}
}

// recordReload records that the given VReg is reloaded in the given block.
// This is used to determine the spill location by tracking the lowest common ancestor of all the blocks that reloads the value.
func ( *vrState[, , ]) ( ,  ) {
	.spilled = true
	var  
	if  := .lca;  ==  {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("\t\tv%d is reloaded in blk%d,\n", .v.ID(), .ID())
		}
		.lca = 
	} else if  !=  {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("\t\tv%d is reloaded in blk%d, lca=%d\n", .v.ID(), .ID(), .lca.ID())
		}
		.lca = .LowestCommonAncestor(, )
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("updated lca=%d\n", .lca.ID())
		}
	}
}

func ( *Allocator[, , ]) ( *state[, , ],  []RealReg,  RegSet,  RealReg) ( RealReg) {
	 = RealRegInvalid
	// First, check if the preferredMask has any allocatable register.
	if  != RealRegInvalid && !.has() && !.regsInUse.has() {
		return 
	}

	var  programCounter
	var  VReg
	for ,  := range  {
		if .has() {
			continue
		}

		 := .regsInUse.get()
		if  == nil {
			// This is not used at this point.
			return 
		}

		// Real registers in use should not be spilled, so we skip them.
		// For example, if the register is used as an argument register, and it might be
		// spilled and not reloaded when it ends up being used as a temporary to pass
		// stack based argument.
		if .v.IsRealReg() {
			continue
		}

		 :=  == 

		// last == -1 means the value won't be used anymore.
		if  := .lastUse;  == RealRegInvalid ||  ||  == -1 || ( != -1 &&  > ) {
			 = 
			 = 
			 = .v
			if  {
				break
			}
		}
	}

	if  == RealRegInvalid {
		panic("not found any allocatable register")
	}

	if wazevoapi.RegAllocLoggingEnabled {
		fmt.Printf("\tspilling v%d when lastUseAt=%d and regsInUse=%s\n", .ID(), , .regsInUse.format(.regInfo))
	}
	.releaseRealReg()
	return 
}

func ( *state[, , ]) ( []RealReg,  RegSet) RealReg {
	for ,  := range  {
		if !.regsInUse.has() && !.has() {
			return 
		}
	}
	return RealRegInvalid
}

func ( *state[, , ]) ( *blockState[, , ]) {
	.regsInUse.range_(func( RealReg,  *vrState[, , ]) {
		.r = RealRegInvalid
	})
	.regsInUse.reset()
	.endRegs.range_(func( RealReg,  *vrState[, , ]) {
		if .lastUseUpdatedAtBlockID == .currentBlockID && .lastUse == programCounterLiveIn {
			.regsInUse.add(, )
			.r = 
		}
	})
}

func resetBlockState[ Instr,  Block[],  Function[, ]]( *blockState[, , ]) {
	.seen = false
	.visited = false
	.endRegs.reset()
	.startRegs.reset()
	.startFromPredIndex = -1
	.liveIns = .liveIns[:0]
}

func ( *blockState[, , ]) ( *RegisterInfo) {
	fmt.Println("\t\tblockState:")
	fmt.Println("\t\t\tstartRegs:", .startRegs.format())
	fmt.Println("\t\t\tendRegs:", .endRegs.format())
	fmt.Println("\t\t\tstartFromPredIndex:", .startFromPredIndex)
	fmt.Println("\t\t\tvisited:", .visited)
}

// DoAllocation performs register allocation on the given Function.
func ( *Allocator[, , ]) ( ) {
	.livenessAnalysis()
	.alloc()
	.determineCalleeSavedRealRegs()
}

func ( *Allocator[, , ]) ( ) {
	.allocatedCalleeSavedRegs = .allocatedCalleeSavedRegs[:0]
	.state.allocatedRegSet.Range(func( RealReg) {
		if .regInfo.CalleeSavedRegisters.has() {
			.allocatedCalleeSavedRegs = append(.allocatedCalleeSavedRegs, .regInfo.RealRegToVReg[])
		}
	})
	.ClobberedRegisters(.allocatedCalleeSavedRegs)
}

func ( *Allocator[, , ]) ( int32) *blockState[, , ] {
	return .blockStates.GetOrAllocate(int())
}

// phiBlk returns the block that defines the given phi value, nil otherwise.
func ( *vrState[, , ]) ()  {
	if .isPhi {
		return .defBlk
	}
	var  
	return 
}

const (
	programCounterLiveIn  = math.MinInt32
	programCounterLiveOut = math.MaxInt32
)

// liveAnalysis constructs Allocator.blockLivenessData.
// The algorithm here is described in https://pfalcon.github.io/ssabook/latest/book-full.pdf Chapter 9.2.
func ( *Allocator[, , ]) ( ) {
	 := &.state

	for  := VRegID(0);  < vRegIDReservedForRealNum; ++ {
		.getOrAllocateVRegState(VReg().SetRealReg(RealReg()))
	}

	var  
	var  
	for  := .PostOrderBlockIteratorBegin();  != ;  = .PostOrderBlockIteratorNext() {
		// We should gather phi value data.
		for ,  := range .BlockParams(, &.vs) {
			 := .getOrAllocateVRegState()
			.isPhi = true
			.defBlk = 
		}

		 := .ID()
		 := .getOrAllocateBlockState()

		.ss = .ss[:0]
		const (
			 = false
			    = true
		)
		 := .Succs()
		for  := 0;  < ; ++ {
			 := .Succ(, )
			if  ==  {
				continue
			}

			 := .ID()
			 := .getOrAllocateBlockState()
			if !.seen { // This means the back edge.
				continue
			}

			for ,  := range .liveIns {
				if .phiBlk() !=  && .spilled !=  { //nolint:gosimple
					// We use .spilled field to store the flag.
					.spilled = 
					.ss = append(.ss, )
				}
			}
		}

		for  := .InstrRevIteratorBegin();  != ;  = .InstrRevIteratorNext() {

			var ,  VReg
			var  bool
			for _,  = range .Defs(&.vs) {
				if !.IsRealReg() {
					 := .getOrAllocateVRegState()
					 = .isPhi
					// Note: We use .spilled field to store the flag.
					.spilled = 
				}
			}
			for _,  = range .Uses(&.vs) {
				if !.IsRealReg() {
					 := .getOrAllocateVRegState()
					// Note: We use .spilled field to store the flag.
					if .spilled !=  { //nolint:gosimple
						.spilled = 
						.ss = append(.ss, )
					}
				}
			}

			if  {
				if .Valid() && .IsRealReg() {
					// If the destination is a phi value, and the source is a real register, this is the beginning of the function.
					.state.argRealRegs = append(.state.argRealRegs, )
				}
			}
		}

		for ,  := range .ss {
			// We use .spilled field to store the flag.
			if .spilled ==  { //nolint:gosimple
				.liveIns = append(.liveIns, )
				.spilled = false
			}
		}

		.seen = true
	}

	 := .LoopNestingForestRoots()
	for  := 0;  < ; ++ {
		 := .LoopNestingForestRoot()
		.loopTreeDFS(, )
	}
}

// loopTreeDFS implements the Algorithm 9.3 in the book in an iterative way.
func ( *Allocator[, , ]) ( ,  ) {
	.blks = .blks[:0]
	.blks = append(.blks, )

	for len(.blks) > 0 {
		 := len(.blks) - 1
		 := .blks[]
		.blks = .blks[:]
		.ss = .ss[:0]
		const (
			    = false
			 = true
		)
		 := .getOrAllocateBlockState(.ID())
		for ,  := range .liveIns {
			if .phiBlk() !=  {
				.ss = append(.ss, )
				// We use .spilled field to store the flag.
				.spilled = 
			}
		}

		var  []*vrState[, , ]
		 := .LoopNestingForestChildren()
		for  := 0;  < ; ++ {
			 := .LoopNestingForestChild(, )
			 := .ID()
			 := .getOrAllocateBlockState()

			if  == 0 {
				 := len(.liveIns)
				for ,  := range .ss {
					// We use .spilled field to store the flag.
					if .spilled ==  { //nolint:gosimple
						.spilled = 
						// TODO: deduplicate, though I don't think it has much impact.
						.liveIns = append(.liveIns, )
					}
				}
				 = .liveIns[:]
			} else {
				// TODO: deduplicate, though I don't think it has much impact.
				.liveIns = append(.liveIns, ...)
			}

			if .LoopHeader() {
				.blks = append(.blks, )
			}
		}

		if  == 0 {
			// If there's no forest child, we haven't cleared the .spilled field at this point.
			for ,  := range .ss {
				.spilled = false
			}
		}
	}
}

// alloc allocates registers for the given function by iterating the blocks in the reverse postorder.
// The algorithm here is derived from the Go compiler's allocator https://github.com/golang/go/blob/release-branch.go1.21/src/cmd/compile/internal/ssa/regalloc.go
// In short, this is a simply linear scan register allocation where each block inherits the register allocation state from
// one of its predecessors. Each block inherits the selected state and starts allocation from there.
// If there's a discrepancy in the end states between predecessors, the adjustments are made to ensure consistency after allocation is done (which we call "fixing merge state").
// The spill instructions (store into the dedicated slots) are inserted after all the allocations and fixing merge states. That is because
// at the point, we all know where the reloads happen, and therefore we can know the best place to spill the values. More precisely,
// the spill happens in the block that is the lowest common ancestor of all the blocks that reloads the value.
//
// All of these logics are almost the same as Go's compiler which has a dedicated description in the source file ^^.
func ( *Allocator[, , ]) ( ) {
	// First we allocate each block in the reverse postorder (at least one predecessor should be allocated for each block).
	var  
	for  := .ReversePostOrderBlockIteratorBegin();  != ;  = .ReversePostOrderBlockIteratorNext() {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("========== allocating blk%d ========\n", .ID())
		}
		if .Entry() {
			.finalizeStartReg(, )
		}
		.allocBlock(, )
	}
	// After the allocation, we all know the start and end state of each block. So we can fix the merge states.
	for  := .ReversePostOrderBlockIteratorBegin();  != ;  = .ReversePostOrderBlockIteratorNext() {
		.fixMergeState(, )
	}
	// Finally, we insert the spill instructions as we know all the places where the reloads happen.
	.scheduleSpills()
}

func ( *Allocator[, , ]) ( *blockState[, , ]) {
	 := .state.currentBlockID
	for ,  := range .liveIns {
		.lastUse = programCounterLiveIn
		.lastUseUpdatedAtBlockID = 
	}
}

func ( *Allocator[, , ]) ( ,  ) {
	 := .ID()
	 := &.state
	 := .getOrAllocateBlockState()
	if .startFromPredIndex > -1 {
		return
	}

	.currentBlockID = 
	.updateLiveInVRState()

	 := .Preds()
	var  *blockState[, , ]
	switch  {
	case 0: // This is the entry block.
	case 1:
		 := .Pred(, 0).ID()
		 = .getOrAllocateBlockState()
		.startFromPredIndex = 0
	default:
		// TODO: there should be some better heuristic to choose the predecessor.
		for  := 0;  < ; ++ {
			 := .Pred(, ).ID()
			if  := .getOrAllocateBlockState(); .visited {
				 = 
				.startFromPredIndex = 
				break
			}
		}
	}
	if  == nil {
		if !.Entry() {
			panic(fmt.Sprintf("BUG: at lease one predecessor should be visited for blk%d", .ID()))
		}
		for ,  := range .argRealRegs {
			.useRealReg(.RealReg(), .getVRegState(.ID()))
		}
		.startFromPredIndex = 0
	} else {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("allocating blk%d starting from blk%d (on index=%d) \n",
				, .Pred(, .startFromPredIndex).ID(), .startFromPredIndex)
		}
		.resetAt()
	}

	.regsInUse.range_(func( RealReg,  *vrState[, , ]) {
		.startRegs.add(, )
	})
	if wazevoapi.RegAllocLoggingEnabled {
		fmt.Printf("finalized start reg for blk%d: %s\n", .ID(), .startRegs.format(.regInfo))
	}
}

func ( *Allocator[, , ]) ( ,  ) {
	 := .ID()
	 := &.state
	 := .getOrAllocateBlockState()
	.currentBlockID = 

	if .startFromPredIndex < 0 {
		panic("BUG: startFromPredIndex should be set in finalizeStartReg prior to allocBlock")
	}

	// Clears the previous state.
	.regsInUse.range_(func( RealReg,  *vrState[, , ]) { .r = RealRegInvalid })
	.regsInUse.reset()
	// Then set the start state.
	.startRegs.range_(func( RealReg,  *vrState[, , ]) { .useRealReg(, ) })

	 := .ss[:0]

	// Update the last use of each VReg.
	.copies = .copies[:0] // Stores the copy instructions.
	var  programCounter
	var  
	for  := .InstrIteratorBegin();  != ;  = .InstrIteratorNext() {
		var  *vrState[, , ]
		for ,  := range .Uses(&.vs) {
			 = .getVRegState(.ID())
			if !.IsRealReg() {
				.lastUse = 
			}
		}

		if .IsCopy() {
			 := .Defs(&.vs)[0]
			.copies = append(.copies, _copy[, , ]{src: , dstID: .ID()})
			 := .RealReg()
			if  != RealRegInvalid {
				if !.isPhi { // TODO: no idea why do we need this.
					.desiredLoc = newDesiredLocReg()
					 = append(, )
				}
			}
		}
		++
	}

	// Mark all live-out values by checking live-in of the successors.
	// While doing so, we also update the desired register values.
	var  
	var  
	for ,  := 0, .Succs();  < ; ++ {
		 = .Succ(, )
		if  ==  {
			continue
		}

		 := .ID()
		 := .getOrAllocateBlockState()
		for ,  := range .liveIns {
			if .phiBlk() !=  {
				.lastUse = programCounterLiveOut
			}
		}

		if .startFromPredIndex > -1 {
			if wazevoapi.RegAllocLoggingEnabled {
				fmt.Printf("blk%d -> blk%d: start_regs: %s\n", , , .startRegs.format(.regInfo))
			}
			.startRegs.range_(func( RealReg,  *vrState[, , ]) {
				.desiredLoc = newDesiredLocReg()
				 = append(, )
			})
			for ,  := range .BlockParams(, &.vs) {
				 := .getVRegState(.ID())
				if .desiredLoc.realReg() == RealRegInvalid {
					.desiredLoc = desiredLocStack
					 = append(, )
				}
			}
		}
	}

	// Propagate the desired register values from the end of the block to the beginning.
	for ,  := range .copies {
		 := .getVRegState(.dstID)
		 := .desiredLoc.realReg()
		 := .src
		if .phiBlk() !=  && .desiredLoc == desiredLocUnspecified {
			.desiredLoc = newDesiredLocReg()
			 = append(, )
		}
	}

	 = 0
	for  := .InstrIteratorBegin();  != ;  = .InstrIteratorNext() {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Println()
		}

		var  RegSet
		 := .reals[:0]

		// Gather the set of registers that will be used in the current instruction.
		 := .Uses(&.vs)
		for ,  := range  {
			if .IsRealReg() {
				 := .RealReg()
				 = .add()
				if .allocatableSet.has() {
					 = append(, )
				}
			} else {
				 := .getVRegState(.ID())
				if  := .r;  != RealRegInvalid {
					 = .add()
				}
			}
		}

		for ,  := range  {
			if !.IsRealReg() {
				 := .getVRegState(.ID())
				 := .lastUse == 
				 := .r

				if  == RealRegInvalid {
					 = .findOrSpillAllocatable(, .regInfo.AllocatableRegisters[.RegType()], ,
						// Prefer the desired register if it's available.
						.desiredLoc.realReg())
					.recordReload(, )
					.ReloadRegisterBefore(.SetRealReg(), )
					.useRealReg(, )
				}
				if wazevoapi.RegAllocLoggingEnabled {
					fmt.Printf("\ttrying to use v%v on %s\n", .ID(), .regInfo.RealRegName())
				}
				.AssignUse(, .SetRealReg())
				 = .add()
				if  {
					if wazevoapi.RegAllocLoggingEnabled {
						fmt.Printf("\tkill v%d with %s\n", .ID(), .regInfo.RealRegName())
					}
					 = append(, )
				}
			}
		}

		 := .IsIndirectCall()
		if .IsCall() ||  {
			 := RealRegInvalid
			if  {
				 = .vs[0].RealReg()
			}
			.releaseCallerSavedRegs()
		}

		for ,  := range  {
			.releaseRealReg()
		}
		.reals = 

		 := .Defs(&.vs)
		switch len() {
		default:
			// Some instructions define multiple values on real registers.
			// E.g. call instructions (following calling convention) / div instruction on x64 that defines both rax and rdx.
			//
			// Note that currently I assume that such instructions define only the pre colored real registers, not the VRegs
			// that require allocations. If we need to support such case, we need to add the logic to handle it here,
			// though is there any such instruction?
			for ,  := range  {
				if !.IsRealReg() {
					panic("BUG: multiple defs should be on real registers")
				}
				 := .RealReg()
				if .regsInUse.has() {
					.releaseRealReg()
				}
				.useRealReg(, .getVRegState(.ID()))
			}
		case 0:
		case 1:
			 := [0]
			 := .getVRegState(.ID())
			if .IsRealReg() {
				 := .RealReg()
				if .allocatableSet.has() {
					if .regsInUse.has() {
						.releaseRealReg()
					}
					.useRealReg(, )
				}
			} else {
				 := .r

				if  := .desiredLoc.realReg();  != RealRegInvalid {
					if  !=  {
						if (.isPhi && .defBlk == ) ||
							// If this is not a phi and it's already assigned a real reg,
							// this value has multiple definitions, hence we cannot assign the desired register.
							(!.regsInUse.has() &&  == RealRegInvalid) {
							// If the phi value is passed via a real register, we force the value to be in the desired register.
							if wazevoapi.RegAllocLoggingEnabled {
								fmt.Printf("\t\tv%d is phi and desiredReg=%s\n", .ID(), .regInfo.RealRegName())
							}
							if  != RealRegInvalid {
								// If the value is already in a different real register, we release it to change the state.
								// Otherwise, multiple registers might have the same values at the end, which results in
								// messing up the merge state reconciliation.
								.releaseRealReg()
							}
							 = 
							.releaseRealReg()
							.useRealReg(, )
						}
					}
				}

				// Allocate a new real register if `def` is not currently assigned one.
				// It can happen when multiple instructions define the same VReg (e.g. const loads).
				if  == RealRegInvalid {
					if .IsCopy() {
						 := .Uses(&.vs)[0].RealReg()
						if .allocatableSet.has() && !.regsInUse.has() {
							 = 
						}
					}
					if  == RealRegInvalid {
						 := .RegType()
						 = .findOrSpillAllocatable(, .regInfo.AllocatableRegisters[], RegSet(0), RealRegInvalid)
					}
					.useRealReg(, )
				}
				 := .SetRealReg()
				.AssignDef()
				if wazevoapi.RegAllocLoggingEnabled {
					fmt.Printf("\tdefining v%d with %s\n", .ID(), .regInfo.RealRegName())
				}
				if .isPhi {
					if .desiredLoc.stack() { // Stack based phi value.
						.StoreRegisterAfter(, )
						// Release the real register as it's not used anymore.
						.releaseRealReg()
					} else {
						// Only the register based phis are necessary to track the defining instructions
						// since the stack-based phis are already having stores inserted ^.
						 := .phiDefInstListPool.Allocate()
						.instr = 
						.next = .phiDefInstList
						.v = 
						.phiDefInstList = 
					}
				} else {
					.defInstr = 
					.defBlk = 
				}
			}
		}
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Println()
		}
		++
	}

	.regsInUse.range_(func( RealReg,  *vrState[, , ]) { .endRegs.add(, ) })

	.visited = true
	if wazevoapi.RegAllocLoggingEnabled {
		.dump(.regInfo)
	}

	// Reset the desired end location.
	for ,  := range  {
		.desiredLoc = desiredLocUnspecified
	}
	.ss = [:0]

	for  := 0;  < .Succs(); ++ {
		 := .Succ(, )
		if  ==  {
			continue
		}
		// If the successor is not visited yet, finalize the start state.
		.finalizeStartReg(, )
	}
}

func ( *Allocator[, , ]) ( RealReg) {
	 := &.state

	for  := RealReg(0);  < 64; ++ {
		if  ==  { // If this is the call indirect, we should not touch the addr register.
			continue
		}
		if  := .regsInUse.get();  != nil {
			if .v.IsRealReg() {
				continue // This is the argument register as it's already used by VReg backed by the corresponding RealReg.
			}
			if !.regInfo.CallerSavedRegisters.has() {
				// If this is not a caller-saved register, it is safe to keep it across the call.
				continue
			}
			.releaseRealReg()
		}
	}
}

func ( *Allocator[, , ]) ( ,  ) {
	 := .Preds()
	if  <= 1 {
		return
	}

	 := &.state

	// Restores the state at the beginning of the block.
	 := .ID()
	 := .getOrAllocateBlockState()
	 := &.startRegs
	var  RegSet
	for ,  := range  {
		if  != nil {
			 = .add(RealReg())
		}
	}

	if wazevoapi.RegAllocLoggingEnabled {
		fmt.Println("fixMergeState", .ID(), ":", .format(.regInfo))
	}

	.currentBlockID = 
	.updateLiveInVRState()

	for  := 0;  < ; ++ {
		if  == .startFromPredIndex {
			continue
		}

		 := .Pred(, )
		 := .getOrAllocateBlockState(.ID())

		.resetAt()

		// Finds the free registers if any.
		,  := VRegInvalid, VRegInvalid
		if  := .findAllocatable(
			.regInfo.AllocatableRegisters[RegTypeInt], ,
		);  != RealRegInvalid {
			 = FromRealReg(, RegTypeInt)
		}
		if  := .findAllocatable(
			.regInfo.AllocatableRegisters[RegTypeFloat], ,
		);  != RealRegInvalid {
			 = FromRealReg(, RegTypeFloat)
		}

		for  := RealReg(0);  < 64; ++ {
			 := .get()
			if  == nil {
				continue
			}

			 := .regsInUse.get()
			if  != nil && .v.ID() == .v.ID() {
				continue
			}

			 := .v.RegType()
			var  VReg
			if  == RegTypeInt {
				 = 
			} else {
				 = 
			}
			.reconcileEdge(, , , , , , )
		}
	}
}

// reconcileEdge reconciles the register state between the current block and the predecessor for the real register `r`.
//
//   - currentVReg is the current VReg value that sits on the register `r`. This can be VRegInvalid if the register is not used at the end of the predecessor.
//   - desiredVReg is the desired VReg value that should be on the register `r`.
//   - freeReg is the temporary register that can be used to swap the values, which may or may not be used.
//   - typ is the register type of the `r`.
func ( *Allocator[, , ]) ( ,
	 RealReg,
	 ,
	,  *vrState[, , ],
	 VReg,
	 RegType,
) {
	 := .v
	 := VRegInvalid
	if  != nil {
		 = .v
	}
	// There are four cases to consider:
	// 1. currentVReg is valid, but desiredVReg is on the stack.
	// 2. Both currentVReg and desiredVReg are valid.
	// 3. Desired is on a different register than `r` and currentReg is not valid.
	// 4. Desired is on the stack and currentReg is not valid.

	 := &.state
	if .Valid() {
		 := .r
		if  == RealRegInvalid {
			// Case 1: currentVReg is valid, but desiredVReg is on the stack.
			if wazevoapi.RegAllocLoggingEnabled {
				fmt.Printf("\t\tv%d is desired to be on %s, but currently on the stack\n",
					.ID(), .regInfo.RealRegName(),
				)
			}
			// We need to move the current value to the stack, and reload the desired value into the register.
			// TODO: we can do better here.
			.StoreRegisterBefore(.SetRealReg(), .LastInstrForInsertion())
			.releaseRealReg()

			.recordReload(, )
			.ReloadRegisterBefore(.SetRealReg(), .LastInstrForInsertion())
			.useRealReg(, )
			return
		} else {
			// Case 2: Both currentVReg and desiredVReg are valid.
			if wazevoapi.RegAllocLoggingEnabled {
				fmt.Printf("\t\tv%d is desired to be on %s, but currently on %s\n",
					.ID(), .regInfo.RealRegName(), .regInfo.RealRegName(),
				)
			}
			// This case, we need to swap the values between the current and desired values.
			.SwapBefore(
				.SetRealReg(),
				.SetRealReg(),
				,
				.LastInstrForInsertion(),
			)
			.allocatedRegSet = .allocatedRegSet.add(.RealReg())
			.releaseRealReg()
			.releaseRealReg()
			.useRealReg(, )
			.useRealReg(, )
			if wazevoapi.RegAllocLoggingEnabled {
				fmt.Printf("\t\tv%d previously on %s moved to %s\n", .ID(), .regInfo.RealRegName(), .regInfo.RealRegName())
			}
		}
	} else {
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("\t\tv%d is desired to be on %s, current not used\n",
				.ID(), .regInfo.RealRegName(),
			)
		}
		if  := .r;  != RealRegInvalid {
			// Case 3: Desired is on a different register than `r` and currentReg is not valid.
			// We simply need to move the desired value to the register.
			.InsertMoveBefore(
				FromRealReg(, ),
				.SetRealReg(),
				.LastInstrForInsertion(),
			)
			.releaseRealReg()
		} else {
			// Case 4: Both currentVReg and desiredVReg are not valid.
			// We simply need to reload the desired value into the register.
			.recordReload(, )
			.ReloadRegisterBefore(.SetRealReg(), .LastInstrForInsertion())
		}
		.useRealReg(, )
	}
}

func ( *Allocator[, , ]) ( ) {
	 := .state.vrStates
	for  := 0;  <= .MaxIDEncountered(); ++ {
		 := .Get()
		if  == nil {
			continue
		}
		if .spilled {
			.scheduleSpill(, )
		}
	}
}

func ( *Allocator[, , ]) ( ,  *vrState[, , ]) {
	 := .v
	// If the value is the phi value, we need to insert a spill after each phi definition.
	if .isPhi {
		for  := .phiDefInstList;  != nil;  = .next {
			.StoreRegisterAfter(.v, .instr)
		}
		return
	}

	 := .lca
	 := .defBlk
	 := RealRegInvalid
	var  
	if  ==  {
		panic(fmt.Sprintf("BUG: definingBlk should not be nil for %s. This is likley a bug in backend lowering logic", .v.String()))
	}
	if  ==  {
		panic(fmt.Sprintf("BUG: pos should not be nil for %s. This is likley a bug in backend lowering logic", .v.String()))
	}

	if wazevoapi.RegAllocLoggingEnabled {
		fmt.Printf("v%d is spilled in blk%d, lca=blk%d\n", .ID(), .ID(), .ID())
	}
	for  !=  {
		 := .getOrAllocateBlockState(.ID())
		for  := RealReg(0);  < 64; ++ {
			if  := .startRegs.get();  != nil && .v ==  {
				 = 
				// Already in the register, so we can place the spill at the beginning of the block.
				break
			}
		}

		if  != RealRegInvalid {
			break
		}

		 = .Idom()
	}

	if  ==  {
		 := .defInstr
		.Defs(&.vs)
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("schedule spill v%d after %v\n", .ID(), )
		}
		.StoreRegisterAfter(.vs[0], )
	} else {
		// Found an ancestor block that holds the value in the register at the beginning of the block.
		// We need to insert a spill before the last use.
		 := .FirstInstr()
		if wazevoapi.RegAllocLoggingEnabled {
			fmt.Printf("schedule spill v%d before %v\n", .ID(), )
		}
		.StoreRegisterAfter(.SetRealReg(), )
	}
}

// Reset resets the allocator's internal state so that it can be reused.
func ( *Allocator[, , ]) () {
	.state.reset()
	.blockStates.Reset()
	.phiDefInstListPool.Reset()
	.vs = .vs[:0]
}