// Copyright (C) 1994-2024 Lua.org, PUC-Rio.
// Copyright 2024 The zb Authors
// SPDX-License-Identifier: MIT

package luacode

import "math"

// expressionDescriptor describes the location of the result of an expression.
//
// Code generation for variables and expressions can be delayed to allow optimizations;
// an expressionDescriptor structure describes a potentially-delayed variable/expression.
// It has a description of its "main" value
// plus a list of conditional jumps that can also produce its value
// (generated by short-circuit operators "and"/"or").
//
// Equivalent to `expDesc` in upstream Lua.
type expressionDescriptor struct {
	kind expressionKind
	// bits is interpreted based on kind.
	bits uint64
	// strval stores the argument of [codeString].
	strval string

	// t is a patch list of "exit when true".
	// [noJump] indicates an empty list.
	t int
	// f is a patch list of "exit when false".
	// [noJump] indicates an empty list.
	f int
}

func newExpressionDescriptor(kind expressionKind) expressionDescriptor {
	return expressionDescriptor{
		kind: kind,
		t:    noJump,
		f:    noJump,
	}
}

// voidExpression returns an [expressionDescriptor] of [expressionKindVoid].
func voidExpression() expressionDescriptor {
	return newExpressionDescriptor(expressionKindVoid)
}

// codeString returns an [expressionDescriptor] for a string literal constant.
//
// Equivalent to `codestring` in upstream Lua.
func codeString(s string) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindStringConstant)
	e.strval = s
	return e
}

// constantTableExpression returns an [expressionDescriptor] for the k'th constant
// in the [Prototype] Constants table.
func constantTableExpression(k int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindConstant)
	e.bits = uint64(k)
	return e
}

// floatConstantExpression returns an [expressionDescriptor] for a numerical floating point constant.
func floatConstantExpression(f float64) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindFloatConstant)
	e.bits = math.Float64bits(f)
	return e
}

// intConstantExpression returns an [expressionDescriptor] for a numerical integer constant.
func intConstantExpression(i int64) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindIntConstant)
	e.bits = uint64(i)
	return e
}

// nonRelocatableExpression returns an [expressionDescriptor] for a value in a fixed register.
func nonRelocatableExpression(ridx registerIndex) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindNonRelocatable)
	e.bits = uint64(ridx)
	return e
}

// localExpression returns an [expressionDescriptor] for a local variable
// given the register index
// and the index in [parser].activeVars relative to [parser].firstLocal.
func localExpression(ridx registerIndex, vidx uint16) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindLocal)
	e.bits = uint64(ridx) | uint64(vidx)<<8
	return e
}

func upvalueExpression(idx upvalueIndex) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindUpvalue)
	e.bits = uint64(idx)
	return e
}

// constLocalExpression returns an [expressionDescriptor] for a compile-time <const> variable
// given an absolute index in [parser].activeVars.
func constLocalExpression(i int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindConstLocal)
	e.bits = uint64(i)
	return e
}

func indexedExpression(table, idx registerIndex) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindIndexed)
	e.bits = uint64(idx) | uint64(table)<<16
	return e
}

func indexedUpvalueExpression(table upvalueIndex, constIndex uint16) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindIndexUpvalue)
	e.bits = uint64(constIndex) | uint64(table)<<16
	return e
}

func indexIntExpression(table registerIndex, i uint16) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindIndexInt)
	e.bits = uint64(i) | uint64(table)<<16
	return e
}

func indexStringExpression(table registerIndex, constIndex uint16) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindIndexString)
	e.bits = uint64(constIndex) | uint64(table)<<16
	return e
}

func jumpExpression(pc int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindJump)
	e.bits = uint64(pc)
	return e
}

func relocatableExpression(pc int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindRelocatable)
	e.bits = uint64(pc)
	return e
}

func callExpression(pc int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindCall)
	e.bits = uint64(pc)
	return e
}

func varargExpression(pc int) expressionDescriptor {
	e := newExpressionDescriptor(expressionKindVararg)
	e.bits = uint64(pc)
	return e
}

// constantToExpression converts a constant to an [expressionDescriptor].
//
// Equivalent to `const2exp` in upstream Lua.
func constantToExpression(v Value) expressionDescriptor {
	if v.IsNil() {
		return newExpressionDescriptor(expressionKindNil)
	}
	if v.IsString() {
		s, _ := v.Unquoted()
		return codeString(s)
	}
	if v.IsInteger() {
		i, _ := v.Int64(OnlyIntegral)
		return intConstantExpression(i)
	}
	if f, ok := v.Float64(); ok {
		return floatConstantExpression(f)
	}
	if b, ok := v.Bool(); ok {
		if b {
			return newExpressionDescriptor(expressionKindTrue)
		} else {
			return newExpressionDescriptor(expressionKindFalse)
		}
	}
	panic("unhandled Value type")
}

func (e expressionDescriptor) hasJumps() bool {
	return e.t != e.f
}

func (e expressionDescriptor) withJumpLists(from expressionDescriptor) expressionDescriptor {
	e.t = from.t
	e.f = from.f
	return e
}

// toValue returns the argument passed to
// [floatConstantExpression], [intConstantExpression], or [codeString]
// as a [Value].
// It also supports values from [newExpressionDescriptor]
// with kinds [expressionKindNil], [expressionKindFalse], or [expressionKindTrue].
func (e expressionDescriptor) toValue() (_ Value, ok bool) {
	if e.hasJumps() {
		return Value{}, false
	}
	switch e.kind {
	case expressionKindNil:
		return Value{}, true
	case expressionKindFalse:
		return BoolValue(false), true
	case expressionKindTrue:
		return BoolValue(true), true
	case expressionKindIntConstant:
		i, _ := e.intConstant()
		return IntegerValue(i), true
	case expressionKindFloatConstant:
		f, _ := e.floatConstant()
		return FloatValue(f), true
	case expressionKindStringConstant:
		return StringValue(e.strval), true
	default:
		return Value{}, false
	}
}

// isNumeral reports whether e
// was created from [floatConstantExpression] or [intConstantExpression]
// and does not have jumps.
func (e expressionDescriptor) isNumeral() bool {
	return !e.hasJumps() && e.kind == expressionKindIntConstant || e.kind == expressionKindFloatConstant
}

// toNumeral returns the argument passed to
// [floatConstantExpression] or [intConstantExpression]
// as a [Value],
// as long as the expression does not have jumps.
//
// Equivalent to `tonumeral` in upstream Lua.
func (e expressionDescriptor) toNumeral() (_ Value, ok bool) {
	if !e.isNumeral() {
		return Value{}, false
	}
	return e.toValue()
}

// toSignedArg converts a numeral (see [expressionDescriptor.isNumeral])
// into a signed argument (see [ToSignedArg]), if possible.
func (e expressionDescriptor) toSignedArg() (arg uint8, isFloat bool, ok bool) {
	var i int64
	switch e.kind {
	case expressionKindIntConstant:
		i, _ = e.intConstant()
	case expressionKindFloatConstant:
		f, _ := e.floatConstant()
		i, ok = FloatToInteger(f, OnlyIntegral)
		if !ok {
			return 0, true, false
		}
		isFloat = true
	default:
		return 0, false, false
	}

	if e.hasJumps() {
		return 0, isFloat, false
	}
	arg, ok = ToSignedArg(i)
	return arg, isFloat, ok
}

// floatConstant returns the argument passed to [floatConstantExpression].
func (e expressionDescriptor) floatConstant() (_ float64, ok bool) {
	if e.kind != expressionKindFloatConstant {
		return 0, false
	}
	return math.Float64frombits(e.bits), true
}

// intConstant returns the argument passed to [intConstantExpression].
func (e expressionDescriptor) intConstant() (_ int64, ok bool) {
	if e.kind != expressionKindIntConstant {
		return 0, false
	}
	return int64(e.bits), true
}

// stringConstant returns the argument passed to [codeString].
func (e expressionDescriptor) stringConstant() (_ string, ok bool) {
	if e.kind != expressionKindStringConstant {
		return "", false
	}
	return e.strval, true
}

// constantIndex returns the index in the [Prototype] Constants table.
// For [expressionKindIndexUpvalue] or [expressionKindIndexString],
// constantIndex returns the table index constant.
func (e expressionDescriptor) constantIndex() int {
	switch e.kind {
	case expressionKindConstant:
		return int(e.bits)
	case expressionKindIndexUpvalue, expressionKindIndexString:
		return int(e.bits & 0xffff)
	default:
		panic("constIndex not supported on expression")
	}
}

func (e expressionDescriptor) register() registerIndex {
	switch e.kind {
	case expressionKindNonRelocatable, expressionKindLocal:
		return registerIndex(e.bits & 0xff)
	default:
		panic("register not supported on expression")
	}
}

// localIndex returns the index in the [parser] activeVars slice
// for a [localExpression].
func (e expressionDescriptor) localIndex(firstLocal int) int {
	if e.kind != expressionKindLocal {
		panic("localIndex on non-local expression")
	}
	return firstLocal + int(e.bits>>8&0xffff)
}

// upvalueIndex returns the upvalue index passed to [upvalueExpression].
func (e expressionDescriptor) upvalueIndex() upvalueIndex {
	if e.kind != expressionKindUpvalue {
		panic("upvalueIndex on non-upvalue expression")
	}
	return upvalueIndex(e.bits)
}

// constLocalIndex returns the absolute index in the [parser] activeVariables slice
// for a [constLocalExpression].
func (e expressionDescriptor) constLocalIndex() int {
	if e.kind != expressionKindConstLocal {
		panic("constLocalIndex on non-<const> expression")
	}
	return int(e.bits)
}

// tableRegister returns the register holding the table in an index expression.
func (e expressionDescriptor) tableRegister() registerIndex {
	switch e.kind {
	case expressionKindIndexed, expressionKindIndexInt, expressionKindIndexString:
		return registerIndex(e.bits >> 16)
	default:
		panic("tableRegister on non-index expression")
	}
}

func (e expressionDescriptor) withTableRegister(r registerIndex) (_ expressionDescriptor, ok bool) {
	switch e.kind {
	case expressionKindIndexed, expressionKindIndexInt, expressionKindIndexString:
		e.bits = e.bits&^(0xff<<16) | uint64(r)<<16
		return e, true
	default:
		return e, false
	}
}

// tableUpvalue returns the table's upvalue index of the [expressionKindIndexUpvalue] expression.
func (e expressionDescriptor) tableUpvalue() upvalueIndex {
	if e.kind != expressionKindIndexUpvalue {
		panic("tableUpvalue on non-upvalue-index expression")
	}
	return upvalueIndex(e.bits >> 16)
}

// indexRegister returns the table index register of the [expressionKindIndexed] expression.
func (e expressionDescriptor) indexRegister() registerIndex {
	if e.kind != expressionKindIndexed {
		panic("indexRegister on non-index expression")
	}
	return registerIndex(e.bits)
}

// indexInt returns the constant integer of the [expressionKindIndexInt] expression.
func (e expressionDescriptor) indexInt() int64 {
	if e.kind != expressionKindIndexInt {
		panic("indexInt on non-index expression")
	}
	return int64(e.bits)
}

// pc returns the index of the expression's instruction
// in the [Prototype] Code slice.
func (e expressionDescriptor) pc() int {
	switch e.kind {
	case expressionKindJump, expressionKindRelocatable, expressionKindCall, expressionKindVararg:
		return int(e.bits)
	default:
		panic("pc not supported on expression")
	}
}

// expressionKind is an enumeration of types of [expressionDescriptor] values.
type expressionKind int

const (
	// expressionKindVoid is an absence of an expression.
	// Used for errors or an empty expression list.
	// Equivalent to `VVOID` in upstream Lua.
	expressionKindVoid expressionKind = iota
	// expressionKindNil is a nil constant.
	// Equivalent to `VNIL` in upstream Lua.
	expressionKindNil
	// expressionKindTrue is a true constant.
	// Equivalent to `VTRUE` in upstream Lua.
	expressionKindTrue
	// expressionKindFalse is a false constant.
	// Equivalent to `VFALSE` in upstream Lua.
	expressionKindFalse
	// expressionKindConstant is a constant stored in the [Prototype] Constants table.
	// Equivalent to `VK` in upstream Lua.
	expressionKindConstant
	// expressionKindFloatConstant is a floating-point constant.
	// Equivalent to `VKFLT` in upstream Lua.
	expressionKindFloatConstant
	// expressionKindIntConstant is an integer constant.
	// Equivalent to `VKINT` in upstream Lua.
	expressionKindIntConstant
	// expressionKindStringConstant is a string constant.
	// Equivalent to `VKSTR` in upstream Lua.
	expressionKindStringConstant
	// expressionKindNonRelocatable is an expression whose value is stored in a fixed register.
	// Equivalent to `VNONRELOC` in upstream Lua.
	expressionKindNonRelocatable
	// expressionKindLocal is a local variable reference.
	// Equivalent to `VLOCAL` in upstream Lua.
	expressionKindLocal
	// upvalue variable; info = index of upvalue in 'upvalues'
	// Equivalent to `VUPVAL` in upstream Lua.
	expressionKindUpvalue
	// expressionKindConstLocal is a local constant variable
	// (i.e. one marked as "<const>").
	// Equivalent to `VCONST` in upstream Lua.
	expressionKindConstLocal
	// expressionKindIndexed is a lookup of a key stored in one register
	// in a table stored in another register.
	// Equivalent to `VINDEXED` in upstream Lua.
	expressionKindIndexed
	// expressionKindIndexUpvalue is a lookup of a key stored in the [Prototype] Constants table
	// in a table stored in an upvalue.
	// Equivalent to `VINDEXUP` in upstream Lua.
	expressionKindIndexUpvalue
	// expressionKindIndexInt is a lookup of an integer key
	// in a table stored in a register.
	// Equivalent to `VINDEXI` in upstream Lua.
	expressionKindIndexInt
	// expressionKindIndexString is a lookup of a string key
	// stored in the [Prototype] Constants table
	// in a table stored in a register.
	// Equivalent to `VINDEXSTR` in upstream Lua.
	expressionKindIndexString
	// expressionKindJump is a test/comparison.
	// Equivalent to `VJMP` in upstream Lua.
	expressionKindJump
	// expressionKindRelocatable whose result can be placed in any register.
	// Equivalent to `VRELOC` in upstream Lua.
	expressionKindRelocatable
	// expressionKindCall is a function call.
	// Equivalent to `VCALL` in upstream Lua.
	expressionKindCall
	// expressionKindVararg is a vararg expression.
	// Equivalent to `VVARARG` in upstream Lua.
	expressionKindVararg
)

// isCompileTimeConstant reports whether an [expressionDescriptor] of this kind
// is a compile-time constant.
func (k expressionKind) isCompileTimeConstant() bool {
	return expressionKindNil <= k && k <= expressionKindStringConstant
}

// isVariable reports whether an [expressionDescriptor] of this kind
// accesses a variable.
//
// Equivalent to `vkisvar` in upstream Lua.
func (k expressionKind) isVariable() bool {
	return expressionKindLocal <= k && k <= expressionKindIndexString
}

// isIndexed reports whether an [expressionDescriptor] of this kind
// is indexing a table.
//
// Equivalent to `vkisvar` in upstream Lua.
func (k expressionKind) isIndexed() bool {
	return expressionKindIndexed <= k && k <= expressionKindIndexString
}

// hasMultipleReturns reports whether an [expressionDescriptor] of this kind
// can return multiple values.
//
// Equivalent to `hasmultret` in upstream Lua.
func (k expressionKind) hasMultipleReturns() bool {
	return k == expressionKindCall || k == expressionKindVararg
}
