-- AVM ISA v1 (Lean-only): Static Type Checker

import SilverSight.AVMIsa.Instr

namespace SilverSight.AVMIsa

/-- Helper to check Primitive type correctness. -/
def checkPrim : Prim → List AvmTy → Option (List AvmTy)
  | .not => fun stack =>
      match stack with
      | .bool :: xs => some (.bool :: xs)
      | _ => none
  | .and | .or => fun stack =>
      match stack with
      | .bool :: .bool :: xs => some (.bool :: xs)
      | _ => none
  | .addSatQ0 | .subSatQ0 => fun stack =>
      match stack with
      | .q0_16 :: .q0_16 :: xs => some (.q0_16 :: xs)
      | _ => none
  | .addSatQ16 | .subSatQ16 | .mulSatQ16 | .divSatQ16 => fun stack =>
      match stack with
      | .q16_16 :: .q16_16 :: xs => some (.q16_16 :: xs)
      | _ => none
  | .ltQ16 | .eqQ16 => fun stack =>
      match stack with
      | .q16_16 :: .q16_16 :: xs => some (.bool :: xs)
      | _ => none

/-- One-step static type checking of a single instruction.
    Returns `some (new_stack, new_locals)` if type-correct, else `none`. -/
def checkInstr (instr : Instr) (stack : List AvmTy) (locals : List (Option AvmTy)) :
    Option (List AvmTy × List (Option AvmTy)) :=
  match instr with
  | .push v => some (v.ty :: stack, locals)
  | .pop =>
      match stack with
      | _ :: xs => some (xs, locals)
      | [] => none
  | .dup =>
      match stack with
      | x :: xs => some (x :: x :: xs, locals)
      | [] => none
  | .swap =>
      match stack with
      | x :: y :: xs => some (y :: x :: xs, locals)
      | _ => none
  | .load i =>
      match locals[i]? with
      | some (some ty) => some (ty :: stack, locals)
      | _ => none
  | .store i =>
      if i < locals.length then
        match stack with
        | x :: xs => some (xs, locals.set i (some x))
        | [] => none
      else none
  | .jump _ => some (stack, locals)
  | .jumpIf _ =>
      match stack with
      | .bool :: xs => some (xs, locals)
      | _ => none
  | .prim p =>
      match checkPrim p stack with
      | some stack' => some (stack', locals)
      | none => none
  | .halt => some (stack, locals)

/-- Check a whole instruction program sequentially. -/
def checkProgram (prog : List Instr) (initStack : List AvmTy) (initLocals : List (Option AvmTy)) : Bool :=
  let rec loop (pc : Nat) (stack : List AvmTy) (locals : List (Option AvmTy)) (fuel : Nat) : Bool :=
    match fuel with
    | 0 => false
    | f + 1 =>
        match prog[pc]? with
        | none => true -- reached end / out of bounds is checked at runtime, statically we halt checking
        | some instr =>
            match instr with
            | .halt => true
            | .jump target =>
                if target < prog.length then loop target stack locals f
                else false
            | .jumpIf target =>
                match stack with
                | .bool :: xs =>
                    if target < prog.length then
                      loop (pc + 1) xs locals f && loop target xs locals f
                    else false
                | _ => false
            | _ =>
                match checkInstr instr stack locals with
                | some (stack', locals') => loop (pc + 1) stack' locals' f
                | none => false
  loop 0 initStack initLocals (prog.length + 1)

end SilverSight.AVMIsa
