(* AVM ISA v1 — Coq Formalization *)
Require Import ZArith List.
Local Open Scope Z_scope.

Module AVM.
  Definition q16_scale : Z := 65536.

  Inductive AvmTy : Set := Q0_16 | Q16_16 | Bool.

  Inductive AvmVal : Set := Vq0 (x : Z) | Vq16 (x : Z) | Vbool (b : bool).

  Inductive Prim : Set :=
    AddQ16 | SubQ16 | MulQ16 | DivQ16 | LtQ16 | EqQ16 | And | Or | Not.

  Inductive Instr : Set :=
    Push_q16 (x : Z) | Push_bool (b : bool) | Pop | Dup | Swap
  | Load (i : nat) | Store (i : nat) | Jump (t : nat) | Jump_if (t : nat)
  | Prim_op (p : Prim) | Halt.

  Record State : Set := mkState {
    pc : nat; stack : list AvmVal; locals : list (option AvmVal); halted : bool
  }.

  Definition empty_state : State := mkState 0 nil nil false.

  Definition get_local (s : State) (i : nat) : option AvmVal :=
    match List.nth_error s.(locals) i with
    | Some (Some v) => Some v | _ => None
    end.

  Definition q16_mul (a b : Z) : Z := Z.div (a * b) q16_scale.
  Definition q16_div (a b : Z) : Z :=
    if Z.eqb b 0 then q16_scale else Z.div (a * q16_scale) b.

  Definition exec_prim (p : Prim) (a b : option AvmVal) : option AvmVal :=
    match p, a, b with
    | AddQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vq16 (x + y))
    | SubQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vq16 (x - y))
    | MulQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vq16 (q16_mul x y))
    | DivQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vq16 (q16_div x y))
    | LtQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vbool (Z.ltb x y))
    | EqQ16, Some (Vq16 x), Some (Vq16 y) => Some (Vbool (Z.eqb x y))
    | And, Some (Vbool x), Some (Vbool y) => Some (Vbool (x && y))
    | Or, Some (Vbool x), Some (Vbool y) => Some (Vbool (x || y))
    | Not, Some (Vbool x), None => Some (Vbool (negb x))
    | _, _, _ => None
    end.

  Definition step (s : State) (prog : list Instr) : option State :=
    if s.(halted) then None else
    match List.nth_error prog s.(pc) with
    | None => Some (mkState s.(pc) s.(stack) s.(locals) true)
    | Some instr =>
      let new_pc := S s.(pc) in
      let push v := Some (mkState new_pc (v :: s.(stack)) s.(locals) false) in
      match instr with
      | Push_q16 x => push (Vq16 x)
      | Push_bool x => push (Vbool x)
      | Pop =>
        match s.(stack) with
        | _ :: rest => Some (mkState new_pc rest s.(locals) false)
        | nil => None
        end
      | Dup =>
        match s.(stack) with
        | v :: _ => push v | nil => None
        end
      | Swap =>
        match s.(stack) with
        | a :: b :: rest => Some (mkState new_pc (b :: a :: rest) s.(locals) false)
        | _ => None
        end
      | Load i =>
        match get_local s i with
        | Some v => push v | None => None
        end
      | Store i =>
        match s.(stack) with
        | v :: rest =>
          Some (mkState new_pc rest
            (List.firstn i s.(locals) ++ Some v :: List.skipn (S i) s.(locals)) false)
        | nil => None
        end
      | Jump t => Some (mkState t s.(stack) s.(locals) false)
      | Jump_if t =>
        match s.(stack) with
        | Vbool true :: rest => Some (mkState t rest s.(locals) false)
        | Vbool false :: rest => Some (mkState new_pc rest s.(locals) false)
        | _ :: _ => None | nil => None
        end
      | Prim_op p =>
        (match p with
        | Not =>
          match s.(stack) with
          | a :: rest =>
            match exec_prim p (Some a) None with
            | Some v => Some (mkState new_pc (v :: rest) s.(locals) false)
            | None => None
            end
          | nil => None
          end
        | _ =>
          match s.(stack) with
          | a :: b :: rest =>
            match exec_prim p (Some a) (Some b) with
            | Some v => Some (mkState new_pc (v :: rest) s.(locals) false)
            | None => None
            end
          | _ => None
          end
        end)
      | Halt => Some (mkState s.(pc) s.(stack) s.(locals) true)
      end
    end.

  Fixpoint run (s : State) (prog : list Instr) (fuel : nat) {struct fuel} : option State :=
    match fuel with
    | O => Some s
    | S k =>
      if s.(halted) then Some s else
      match step s prog with
      | None => None | Some s' => run s' prog k
      end
    end.

  Lemma step_halted (s : State) (prog : list Instr) (h : s.(halted) = true) :
    step s prog = None.
  Proof. unfold step; rewrite h; reflexivity. Qed.

End AVM.