ReefVM/src/programmatic.ts

import type { Bytecode, Constant } from "./bytecode"
import { OpCode } from "./opcode"
import { toValue } from "./value"

// Instruction types
type PrimitiveValue = number | string | boolean | null

type InstructionTuple =
  // Stack
  | ["PUSH", PrimitiveValue]
  | ["POP"]
  | ["DUP"]

  // Variables
  | ["LOAD", string]
  | ["STORE", string]

  // Arithmetic
  | ["ADD"] | ["SUB"] | ["MUL"] | ["DIV"] | ["MOD"]

  // Comparison
  | ["EQ"] | ["NEQ"] | ["LT"] | ["GT"] | ["LTE"] | ["GTE"]

  // Logical
  | ["NOT"]

  // Control flow
  | ["JUMP", string | number]
  | ["JUMP_IF_FALSE", string | number]
  | ["JUMP_IF_TRUE", string | number]
  | ["BREAK"]

  // Exception handling
  | ["PUSH_TRY", string | number]
  | ["PUSH_FINALLY", string | number]
  | ["POP_TRY"]
  | ["THROW"]

  // Functions
  | ["MAKE_FUNCTION", string[], string | number]
  | ["CALL"]
  | ["TAIL_CALL"]
  | ["RETURN"]

  // Arrays
  | ["MAKE_ARRAY", number]
  | ["ARRAY_GET"]
  | ["ARRAY_SET"]
  | ["ARRAY_PUSH"]
  | ["ARRAY_LEN"]

  // Dicts
  | ["MAKE_DICT", number]
  | ["DICT_GET"]
  | ["DICT_SET"]
  | ["DICT_HAS"]

  // Native
  | ["CALL_NATIVE", string]
  
  // Special
  | ["HALT"]

type LabelDefinition = [string] // Just ".label_name"

type ProgramItem = InstructionTuple | LabelDefinition

function isLabelDefinition(item: ProgramItem): item is LabelDefinition {
  return item.length === 1 && typeof item[0] === "string" && item[0].startsWith(".")
}

function isLabelReference(value: string | number): value is string {
  return typeof value === "string" && value.startsWith(".")
}

function parseFunctionParams(params: string[]): {
  params: string[]
  defaults: Record<string, number>
  variadic: boolean
  named: boolean
  defaultConstants: Constant[]
} {
  const resultParams: string[] = []
  const defaults: Record<string, number> = {}
  const defaultConstants: Constant[] = []
  let variadic = false
  let named = false

  for (const param of params) {
    if (param.startsWith("@")) {
      // Named parameter
      named = true
      resultParams.push(param.slice(1))
    } else if (param.startsWith("...")) {
      // Variadic parameter
      variadic = true
      resultParams.push(param.slice(3))
    } else if (param.includes("=")) {
      // Default parameter
      const [name, defaultValue] = param.split("=").map(s => s.trim())
      resultParams.push(name!)

      // Parse default value
      if (/^-?\d+(\.\d+)?$/.test(defaultValue!)) {
        defaultConstants.push(toValue(parseFloat(defaultValue!)))
      } else if (defaultValue === "true") {
        defaultConstants.push(toValue(true))
      } else if (defaultValue === "false") {
        defaultConstants.push(toValue(false))
      } else if (defaultValue === "null") {
        defaultConstants.push(toValue(null))
      } else if (/^['"].*['"]$/.test(defaultValue!)) {
        defaultConstants.push(toValue(defaultValue!.slice(1, -1)))
      } else {
        throw new Error(`Invalid default value: ${defaultValue}`)
      }

      defaults[name!] = -1 // Will be fixed after we know constant indices
    } else {
      // Regular parameter
      resultParams.push(param)
    }
  }

  return { params: resultParams, defaults, variadic, named, defaultConstants }
}

export function fromInstructions(program: ProgramItem[]): Bytecode {
  const constants: Constant[] = []
  const instructions: any[] = []
  const labels = new Map<string, number>()

  // First pass: collect labels and their positions
  const filteredProgram: InstructionTuple[] = []
  for (const item of program) {
    if (isLabelDefinition(item)) {
      const labelName = item[0].slice(1) // Remove leading "."
      labels.set(labelName, filteredProgram.length)
    } else {
      filteredProgram.push(item as InstructionTuple)
    }
  }

  // Second pass: build instructions
  for (let i = 0; i < filteredProgram.length; i++) {
    const item = filteredProgram[i]!
    const op = item[0] as string
    const opCode = OpCode[op as keyof typeof OpCode]

    if (opCode === undefined) {
      throw new Error(`Unknown opcode: ${op}`)
    }

    let operandValue: number | string | undefined = undefined

    if (item.length > 1) {
      const operand = item[1]

      switch (op) {
        case "PUSH":
          // Add to constants pool
          constants.push(toValue(operand as PrimitiveValue))
          operandValue = constants.length - 1
          break

        case "MAKE_FUNCTION": {
          const params = operand as string[]
          const body = item[2]

          if (body === undefined) {
            throw new Error("MAKE_FUNCTION requires body address")
          }

          const { params: resultParams, defaults, variadic, named, defaultConstants } = parseFunctionParams(params)

          // Add default constants to pool and update indices
          const defaultIndices: Record<string, number> = {}
          for (const [paramName, _] of Object.entries(defaults)) {
            const defaultConst = defaultConstants.shift()!
            constants.push(defaultConst)
            defaultIndices[paramName] = constants.length - 1
          }

          // Resolve body label or use numeric value
          let bodyAddress: number
          if (isLabelReference(body)) {
            const labelName = body.slice(1)
            const labelPos = labels.get(labelName)
            if (labelPos === undefined) {
              throw new Error(`Undefined label: ${labelName}`)
            }
            bodyAddress = labelPos
          } else {
            bodyAddress = body as number
          }

          // Add function definition to constants
          constants.push({
            type: "function_def",
            params: resultParams,
            defaults: defaultIndices,
            body: bodyAddress,
            variadic,
            named
          })

          operandValue = constants.length - 1
          break
        }

        case "JUMP":
        case "JUMP_IF_FALSE":
        case "JUMP_IF_TRUE": {
          // Relative jump
          if (isLabelReference(operand as string | number)) {
            const labelName = (operand as string).slice(1)
            const labelPos = labels.get(labelName)
            if (labelPos === undefined) {
              throw new Error(`Undefined label: ${labelName}`)
            }
            operandValue = labelPos - (i + 1) // Relative offset
          } else {
            operandValue = operand as number
          }
          break
        }

        case "PUSH_TRY":
        case "PUSH_FINALLY": {
          // Absolute address
          if (isLabelReference(operand as string | number)) {
            const labelName = (operand as string).slice(1)
            const labelPos = labels.get(labelName)
            if (labelPos === undefined) {
              throw new Error(`Undefined label: ${labelName}`)
            }
            operandValue = labelPos
          } else {
            operandValue = operand as number
          }
          break
        }

        case "LOAD":
        case "STORE":
        case "CALL_NATIVE":
          // String operand
          operandValue = operand as string
          break

        case "MAKE_ARRAY":
        case "MAKE_DICT":
          // Numeric operand
          operandValue = operand as number
          break

        default:
          throw new Error(`Unexpected operand for ${op}`)
      }
    }

    instructions.push({
      op: opCode,
      operand: operandValue
    })
  }

  // Build labels map for debugger (instruction index -> label name)
  const labelsByIndex = new Map<number, string>()
  for (const [name, index] of labels.entries()) {
    labelsByIndex.set(index, name)
  }

  return {
    instructions,
    constants,
    labels: labelsByIndex
  }
}