github.com/wasilibs/wazerox@v0.0.0-20240124024944-4923be63ab5f/internal/engine/compiler/compiler_value_location.go

package compiler

import (
	"fmt"
	"strings"

	"github.com/wasilibs/wazerox/internal/asm"
	"github.com/wasilibs/wazerox/internal/wasm"
)

var (
	// unreservedGeneralPurposeRegisters contains unreserved general purpose registers of integer type.
	unreservedGeneralPurposeRegisters []asm.Register

	// unreservedVectorRegisters contains unreserved vector registers.
	unreservedVectorRegisters []asm.Register
)

func isGeneralPurposeRegister(r asm.Register) bool {
	return unreservedGeneralPurposeRegisters[0] <= r && r <= unreservedGeneralPurposeRegisters[len(unreservedGeneralPurposeRegisters)-1]
}

func isVectorRegister(r asm.Register) bool {
	return unreservedVectorRegisters[0] <= r && r <= unreservedVectorRegisters[len(unreservedVectorRegisters)-1]
}

// runtimeValueLocation corresponds to each variable pushed onto the wazeroir (virtual) stack,
// and it has the information about where it exists in the physical machine.
// It might exist in registers, or maybe on in the non-virtual physical stack allocated in memory.
type runtimeValueLocation struct {
	valueType runtimeValueType
	// register is set to asm.NilRegister if the value is stored in the memory stack.
	register asm.Register
	// conditionalRegister is set to conditionalRegisterStateUnset if the value is not on the conditional register.
	conditionalRegister asm.ConditionalRegisterState
	// stackPointer is the location of this value in the memory stack at runtime,
	stackPointer uint64
}

func (v *runtimeValueLocation) getRegisterType() (ret registerType) {
	switch v.valueType {
	case runtimeValueTypeI32, runtimeValueTypeI64:
		ret = registerTypeGeneralPurpose
	case runtimeValueTypeF32, runtimeValueTypeF64,
		runtimeValueTypeV128Lo, runtimeValueTypeV128Hi:
		ret = registerTypeVector
	default:
		panic("BUG")
	}
	return
}

type runtimeValueType byte

const (
	runtimeValueTypeNone runtimeValueType = iota
	runtimeValueTypeI32
	runtimeValueTypeI64
	runtimeValueTypeF32
	runtimeValueTypeF64
	runtimeValueTypeV128Lo
	runtimeValueTypeV128Hi
)

func (r runtimeValueType) String() (ret string) {
	switch r {
	case runtimeValueTypeI32:
		ret = "i32"
	case runtimeValueTypeI64:
		ret = "i64"
	case runtimeValueTypeF32:
		ret = "f32"
	case runtimeValueTypeF64:
		ret = "f64"
	case runtimeValueTypeV128Lo:
		ret = "v128.lo"
	case runtimeValueTypeV128Hi:
		ret = "v128.hi"
	}
	return
}

func (v *runtimeValueLocation) setRegister(reg asm.Register) {
	v.register = reg
	v.conditionalRegister = asm.ConditionalRegisterStateUnset
}

func (v *runtimeValueLocation) onRegister() bool {
	return v.register != asm.NilRegister && v.conditionalRegister == asm.ConditionalRegisterStateUnset
}

func (v *runtimeValueLocation) onStack() bool {
	return v.register == asm.NilRegister && v.conditionalRegister == asm.ConditionalRegisterStateUnset
}

func (v *runtimeValueLocation) onConditionalRegister() bool {
	return v.conditionalRegister != asm.ConditionalRegisterStateUnset
}

func (v *runtimeValueLocation) String() string {
	var location string
	if v.onStack() {
		location = fmt.Sprintf("stack(%d)", v.stackPointer)
	} else if v.onConditionalRegister() {
		location = fmt.Sprintf("conditional(%d)", v.conditionalRegister)
	} else if v.onRegister() {
		location = fmt.Sprintf("register(%s)", registerNameFn(v.register))
	}
	return fmt.Sprintf("{type=%s,location=%s}", v.valueType, location)
}

func newRuntimeValueLocationStack() runtimeValueLocationStack {
	return runtimeValueLocationStack{
		unreservedVectorRegisters:         unreservedVectorRegisters,
		unreservedGeneralPurposeRegisters: unreservedGeneralPurposeRegisters,
	}
}

// runtimeValueLocationStack represents the wazeroir virtual stack
// where each item holds the location information about where it exists
// on the physical machine at runtime.
//
// Notably this is only used in the compilation phase, not runtime,
// and we change the state of this struct at every wazeroir operation we compile.
// In this way, we can see where the operands of an operation (for example,
// two variables for wazeroir add operation.) exist and check the necessity for
// moving the variable to registers to perform actual CPU instruction
// to achieve wazeroir's add operation.
type runtimeValueLocationStack struct {
	// stack holds all the variables.
	stack []runtimeValueLocation
	// sp is the current stack pointer.
	sp uint64
	// usedRegisters is the bit map to track the used registers.
	usedRegisters usedRegistersMask
	// stackPointerCeil tracks max(.sp) across the lifespan of this struct.
	stackPointerCeil uint64
	// unreservedGeneralPurposeRegisters and unreservedVectorRegisters hold
	// architecture dependent unreserved register list.
	unreservedGeneralPurposeRegisters, unreservedVectorRegisters []asm.Register
}

func (v *runtimeValueLocationStack) reset() {
	stack := v.stack[:0]
	*v = runtimeValueLocationStack{
		unreservedVectorRegisters:         unreservedVectorRegisters,
		unreservedGeneralPurposeRegisters: unreservedGeneralPurposeRegisters,
		stack:                             stack,
	}
}

func (v *runtimeValueLocationStack) String() string {
	var stackStr []string
	for i := uint64(0); i < v.sp; i++ {
		stackStr = append(stackStr, v.stack[i].String())
	}
	usedRegisters := v.usedRegisters.list()
	return fmt.Sprintf("sp=%d, stack=[%s], used_registers=[%s]", v.sp, strings.Join(stackStr, ","), strings.Join(usedRegisters, ","))
}

// cloneFrom clones the values on `from` into self except for the slice of .stack field.
// The content on .stack will be copied from the origin to self, and grow the underlying slice
// if necessary.
func (v *runtimeValueLocationStack) cloneFrom(from runtimeValueLocationStack) {
	// Assigns the same values for fields except for the stack which we want to reuse.
	prev := v.stack
	*v = from
	v.stack = prev[:cap(prev)] // Expand the length to the capacity so that we can minimize "diff" below.
	// Copy the content in the stack.
	if diff := int(from.sp) - len(v.stack); diff > 0 {
		v.stack = append(v.stack, make([]runtimeValueLocation, diff)...)
	}
	copy(v.stack, from.stack[:from.sp])
}

// pushRuntimeValueLocationOnRegister creates a new runtimeValueLocation with a given register and pushes onto
// the location stack.
func (v *runtimeValueLocationStack) pushRuntimeValueLocationOnRegister(reg asm.Register, vt runtimeValueType) (loc *runtimeValueLocation) {
	loc = v.push(reg, asm.ConditionalRegisterStateUnset)
	loc.valueType = vt
	return
}

// pushRuntimeValueLocationOnRegister creates a new runtimeValueLocation and pushes onto the location stack.
func (v *runtimeValueLocationStack) pushRuntimeValueLocationOnStack() (loc *runtimeValueLocation) {
	loc = v.push(asm.NilRegister, asm.ConditionalRegisterStateUnset)
	loc.valueType = runtimeValueTypeNone
	return
}

// pushRuntimeValueLocationOnRegister creates a new runtimeValueLocation with a given conditional register state
// and pushes onto the location stack.
func (v *runtimeValueLocationStack) pushRuntimeValueLocationOnConditionalRegister(state asm.ConditionalRegisterState) (loc *runtimeValueLocation) {
	loc = v.push(asm.NilRegister, state)
	loc.valueType = runtimeValueTypeI32
	return
}

// push a runtimeValueLocation onto the stack.
func (v *runtimeValueLocationStack) push(reg asm.Register, conditionalRegister asm.ConditionalRegisterState) (ret *runtimeValueLocation) {
	if v.sp >= uint64(len(v.stack)) {
		// This case we need to grow the stack capacity by appending the item,
		// rather than indexing.
		v.stack = append(v.stack, runtimeValueLocation{})
	}
	ret = &v.stack[v.sp]
	ret.register, ret.conditionalRegister, ret.stackPointer = reg, conditionalRegister, v.sp
	v.sp++
	// stackPointerCeil must be set after sp is incremented since
	// we skip the stack grow if len(stack) >= basePointer+stackPointerCeil.
	if v.sp > v.stackPointerCeil {
		v.stackPointerCeil = v.sp
	}
	return
}

func (v *runtimeValueLocationStack) pop() (loc *runtimeValueLocation) {
	v.sp--
	loc = &v.stack[v.sp]
	return
}

func (v *runtimeValueLocationStack) popV128() (loc *runtimeValueLocation) {
	v.sp -= 2
	loc = &v.stack[v.sp]
	return
}

func (v *runtimeValueLocationStack) peek() (loc *runtimeValueLocation) {
	loc = &v.stack[v.sp-1]
	return
}

func (v *runtimeValueLocationStack) releaseRegister(loc *runtimeValueLocation) {
	v.markRegisterUnused(loc.register)
	loc.register = asm.NilRegister
	loc.conditionalRegister = asm.ConditionalRegisterStateUnset
}

func (v *runtimeValueLocationStack) markRegisterUnused(regs ...asm.Register) {
	for _, reg := range regs {
		v.usedRegisters.remove(reg)
	}
}

func (v *runtimeValueLocationStack) markRegisterUsed(regs ...asm.Register) {
	for _, reg := range regs {
		v.usedRegisters.add(reg)
	}
}

type registerType byte

const (
	registerTypeGeneralPurpose registerType = iota
	// registerTypeVector represents a vector register which can be used for either scalar float
	// operation or SIMD vector operation depending on the instruction by which the register is used.
	//
	// Note: In normal assembly language, scalar float and vector register have different notations as
	// Vn is for vectors and Qn is for scalar floats on arm64 for example. But on physical hardware,
	// they are placed on the same locations. (Qn means the lower 64-bit of Vn vector register on arm64).
	//
	// In wazero, for the sake of simplicity in the register allocation, we intentionally conflate these two types
	// and delegate the decision to the assembler which is aware of the instruction types for which these registers are used.
	registerTypeVector
)

func (tp registerType) String() (ret string) {
	switch tp {
	case registerTypeGeneralPurpose:
		ret = "int"
	case registerTypeVector:
		ret = "vector"
	}
	return
}

// takeFreeRegister searches for unused registers. Any found are marked used and returned.
func (v *runtimeValueLocationStack) takeFreeRegister(tp registerType) (reg asm.Register, found bool) {
	var targetRegs []asm.Register
	switch tp {
	case registerTypeVector:
		targetRegs = v.unreservedVectorRegisters
	case registerTypeGeneralPurpose:
		targetRegs = v.unreservedGeneralPurposeRegisters
	}
	for _, candidate := range targetRegs {
		if v.usedRegisters.exist(candidate) {
			continue
		}
		return candidate, true
	}
	return 0, false
}

// Search through the stack, and steal the register from the last used
// variable on the stack.
func (v *runtimeValueLocationStack) takeStealTargetFromUsedRegister(tp registerType) (*runtimeValueLocation, bool) {
	for i := uint64(0); i < v.sp; i++ {
		loc := &v.stack[i]
		if loc.onRegister() {
			switch tp {
			case registerTypeVector:
				if loc.valueType == runtimeValueTypeV128Hi {
					panic("BUG: V128Hi must be above the corresponding V128Lo")
				}
				if isVectorRegister(loc.register) {
					return loc, true
				}
			case registerTypeGeneralPurpose:
				if isGeneralPurposeRegister(loc.register) {
					return loc, true
				}
			}
		}
	}
	return nil, false
}

// init sets up the runtimeValueLocationStack which reflects the state of
// the stack at the beginning of the function.
//
// See the diagram in callEngine.stack.
func (v *runtimeValueLocationStack) init(sig *wasm.FunctionType) {
	for _, t := range sig.Params {
		loc := v.pushRuntimeValueLocationOnStack()
		switch t {
		case wasm.ValueTypeI32:
			loc.valueType = runtimeValueTypeI32
		case wasm.ValueTypeI64, wasm.ValueTypeFuncref, wasm.ValueTypeExternref:
			loc.valueType = runtimeValueTypeI64
		case wasm.ValueTypeF32:
			loc.valueType = runtimeValueTypeF32
		case wasm.ValueTypeF64:
			loc.valueType = runtimeValueTypeF64
		case wasm.ValueTypeV128:
			loc.valueType = runtimeValueTypeV128Lo
			hi := v.pushRuntimeValueLocationOnStack()
			hi.valueType = runtimeValueTypeV128Hi
		default:
			panic("BUG")
		}
	}

	// If the len(results) > len(args), the slots for all results are reserved after
	// arguments, so we reflect that into the location stack.
	for i := 0; i < sig.ResultNumInUint64-sig.ParamNumInUint64; i++ {
		_ = v.pushRuntimeValueLocationOnStack()
	}

	// Then push the control frame fields.
	for i := 0; i < callFrameDataSizeInUint64; i++ {
		loc := v.pushRuntimeValueLocationOnStack()
		loc.valueType = runtimeValueTypeI64
	}
}

// getCallFrameLocations returns each field of callFrame's runtime location.
//
// See the diagram in callEngine.stack.
func (v *runtimeValueLocationStack) getCallFrameLocations(sig *wasm.FunctionType) (
	returnAddress, callerStackBasePointerInBytes, callerFunction *runtimeValueLocation,
) {
	offset := callFrameOffset(sig)
	return &v.stack[offset], &v.stack[offset+1], &v.stack[offset+2]
}

// pushCallFrame pushes a call frame's runtime locations onto the stack assuming that
// the function call parameters are already pushed there.
//
// See the diagram in callEngine.stack.
func (v *runtimeValueLocationStack) pushCallFrame(callTargetFunctionType *wasm.FunctionType) (
	returnAddress, callerStackBasePointerInBytes, callerFunction *runtimeValueLocation,
) {
	// If len(results) > len(args), we reserve the slots for the results below the call frame.
	reservedSlotsBeforeCallFrame := callTargetFunctionType.ResultNumInUint64 - callTargetFunctionType.ParamNumInUint64
	for i := 0; i < reservedSlotsBeforeCallFrame; i++ {
		v.pushRuntimeValueLocationOnStack()
	}

	// Push the runtime location for each field of callFrame struct. Note that each of them has
	// uint64 type, and therefore must be treated as runtimeValueTypeI64.

	// callFrame.returnAddress
	returnAddress = v.pushRuntimeValueLocationOnStack()
	returnAddress.valueType = runtimeValueTypeI64
	// callFrame.returnStackBasePointerInBytes
	callerStackBasePointerInBytes = v.pushRuntimeValueLocationOnStack()
	callerStackBasePointerInBytes.valueType = runtimeValueTypeI64
	// callFrame.function
	callerFunction = v.pushRuntimeValueLocationOnStack()
	callerFunction.valueType = runtimeValueTypeI64
	return
}

// usedRegistersMask tracks the used registers in its bits.
type usedRegistersMask uint64

// add adds the given `r` to the mask.
func (u *usedRegistersMask) add(r asm.Register) {
	*u = *u | (1 << registerMaskShift(r))
}

// remove drops the given `r` from the mask.
func (u *usedRegistersMask) remove(r asm.Register) {
	*u = *u & ^(1 << registerMaskShift(r))
}

// exist returns true if the given `r` is used.
func (u *usedRegistersMask) exist(r asm.Register) bool {
	shift := registerMaskShift(r)
	return (*u & (1 << shift)) > 0
}

// list returns the list of debug string of used registers.
// Only used for debugging and testing.
func (u *usedRegistersMask) list() (ret []string) {
	mask := *u
	for i := 0; i < 64; i++ {
		if mask&(1<<i) > 0 {
			ret = append(ret, registerNameFn(registerFromMaskShift(i)))
		}
	}
	return
}