open HolKernel Parse boolLib bossLib

open uvmMemoryTheory
open uvmIRTheory
open uvmValuesTheory
open uvmErrorsTheory
open sumMonadTheory
open DefineUtils

open alistTheory
open combinTheory
open listTheory
open numpairTheory
open optionTheory
open pairTheory
open pred_setTheory
open sptreeTheory
open lcsymtacs
open monadsyntax

val _ = new_theory "uvmInstructionSemantics"
val _ = add_monadsyntax()
val _ = reveal "C" (* The C (flip) combinator is used in this theory *)

(* reg_reader is a reader monad over a memory region. It provides syntax sugar
   for operations on a thread's registers. *)
val _ = type_abbrev("reg_reader", ``:(num -> value or_error) -> α or_error``)
val _ = define_monad("reg_reader", (fn a => ``:^a reg_reader``), 
  ``λ(a : α). (K : α or_error -> α reg_reader) (OK a : α or_error)``,
  ``λ(x : α reg_reader) (fn : α -> β reg_reader) (m : num -> value or_error).
      join (fn <$> x m <*> OK m)``)

val _ = type_abbrev("ssa_map", ``:ssavar -> (num # uvm_type list) list``)

val get_def = Define`
  get (ssa : ssa_map) (v : ssavar) : value reg_reader list =
    MAP (λr m. m r) (MAP FST (ssa v))
`

val getc_def = Define`
    getc ssa (Var v) = get ssa v
  ∧ getc _ (Const _ hd tl) = MAP return (hd::tl)
`

val run_reg_reader_def = Define`
  run_reg_reader (r : α reg_reader) (regs : mem_region) : α or_error =
    r (λn. FST <$> expect (lookup n regs) (InvalidState "undefined register"))
`

val _ = Datatype`
  jumps = <|
    nor_jump : num # (num + value) list;
    exc_jump : (num # (num + value) list) option
  |>
`

(* A _command_ is a message for a μVM-specific action that is outside the
   RISC-V/WMM memory model, such as memory allocation or thread spawning.
   Commands have no ordering relationship to stores; they are evaluated
   immediately and synchronously.
*)
val _ = Datatype`
  command =
  | DoPushStack func_name funcsig (value list) (jumps option)
  | DoPopStack (value resume_values)
  | DoJump num (value list)
  | DoHeapAlloc uvm_type (num option) num
  | DoStackAlloc uvm_type (num option) num
  (* TODO: DoNewStack *)
  | DoGetStack num
  (* TODO: DoSwapStack *)
  | DoNewThread stack_id (value resume_values) num
  | DoThreadExit
`

(* An _operation_ is an atomic instruction in the WMM model. It includes the
   first five operations from the WMM paper (Nm, Ld, St, Com, Rec), along with
   a μVM-specific Cmd operation for actions not defined in WMM (e.g., spawning
   threads).
   
   Background operations, such as DeqSb, are not defined here; they can be
   inserted between foreground operations by the scheduler layer.
*)
val _ = Datatype`
  opn =
  (* Non-memory execution. The reader monad returns an association list of
     newly-assigned register values. *)
  | Nm ((num, value) alist reg_reader)
  (* Load execution. Loads an IRef from the register in the first argument, then
     puts the value at that IRef in the register specified in the second
     argument. See pg 6 of the WMM paper for the semantics of a load. *)
  | Ld (iaddr reg_reader) num
  (* Store execution. Writes the return value of the reader monad to the store
     buffer (sb). Deviates from the WMM spec by NOT clearing the invalidation
     buffer (ib); instead, the scheduler layer may remove ib entries
     arbitrarily. *)
  | St ((iaddr # value) reg_reader)
  (* Commit fence. Blocks unless sb is empty. *)
  | Com
  (* Reconcile fence. Clears the ib. *)
  | Rec
  (* Command execution. Sends a command message to the scheduler layer. *)
  | Cmd (command reg_reader)
`

(* An _operation block_ is a basic block of operations. It is the result of
   compiling the instructions in a μVM bblock into operations and expanding
   any composite-typed parameters into their scalar components.
*)
val _ = Datatype`
  opn_block = <|
    params : (uvm_type list # num list) list;
    has_exc_param : bool; (* register = LENGTH params, always ref<void> *)
    opns : opn list
  |>
`

val _ = type_abbrev("program", ``:func_name |-> opn_block spt``)

val eval_expr_def = Define`
    eval_expr (ssa : ssa_map) (Id _ v : expression) : value list reg_reader =
      sequence (getc ssa v)
  ∧ eval_expr ssa (BinOp op _ v1 v2) = do
      v1 <- HD (getc ssa v1); v2 <- HD (getc ssa v2);
      op <- case op of
            | ADD => return $+
            | SUB => return $-
            | MUL => return $*
            | _ => K (Fail (NotImplemented "most binops"));
      C CONS [] <$> K (int_binop op v1 v2)
    od
  ∧ eval_expr ssa (CmpOp op _ v1 v2) = (
      if (op = EQ ∨ op = NE) then do
        v1 <- sequence (getc ssa v1);
        v2 <- sequence (getc ssa v2);
        return [IntV 1 (if (v1 = v2) ⇔ (op = EQ) then 1 else 0)]
      od else do
        v1 <- HD (getc ssa v1);
        v2 <- HD (getc ssa v2);
        c <- K (compare_unsigned_values v1 v2);
        b <- case op of
             | UGE => return (c = Equal ∨ c = Greater)
             | UGT => return (c = Greater)
             | ULE => return (c = Equal ∨ c = Less)
             | ULT => return (c = Less)
             | _ => K (Fail (NotImplemented "most cmpops"));
        return [IntV 1 (if b then 1 else 0)]
      od)
  ∧ eval_expr ssa (ConvOp REFCAST ty1 ty2 v) =
     (case ty1, ty2 of
      | Ref (Type Void), Ref (Type t) => (
          HD (getc ssa v) >>= λv. case v of
          | RefV _ a => return [RefV t a]
          | _ => K (Fail (TypeMismatch "REFCAST" ty1 (type_of_value v))))
      | Ref (Type t), Ref (Type Void) => (
          HD (getc ssa v) >>= λv. case v of
          | RefV _ a => return [RefV t a]
          | _ => K (Fail (TypeMismatch "REFCAST" ty1 (type_of_value v))))
      | _ => K (Fail (TypeMismatch "REFCAST" ty1 ty2)))
  ∧ eval_expr _ (ConvOp _ _ _ _) =
      K (Fail (NotImplemented "non-REFCAST conversions"))
    (* TODO: Select *)
  ∧ eval_expr _ (Select _ _ _ _ _) =
      K (Fail (NotImplemented "SELECT"))
  ∧ eval_expr ssa (ExtractValue ty off v) =
     (let members = member_types (SOME off) ty in
      let sizeof = LENGTH o component_types NONE in do
        K (assert (off < LENGTH members) (ReadOutOfBounds ty off));
        vs <- sequence (getc ssa v);
        let out = TAKE (sizeof (EL off members))
                       (DROP (SUM (MAP sizeof (TAKE off members))) vs) in
        K (assert (¬NULL out) (ReadOutOfBounds ty off)) *>
        return out
      od)
  ∧ eval_expr ssa (InsertValue ty off base el) =
     (let members = member_types (SOME off) ty in
      let sizeof = LENGTH o component_types NONE in do
        K (assert (off < LENGTH members) (WriteOutOfBounds ty off));
        vs <- sequence (getc ssa base);
        vs' <- sequence (getc ssa el);
        let out =
          (TAKE (SUM (MAP sizeof (TAKE off members))) vs) ++ vs' ++
          (DROP (SUM (MAP sizeof (TAKE (SUC off) members))) vs) in
        K (assert (LENGTH vs = LENGTH out) (WriteOutOfBounds ty off)) *>
        return out
      od)(*
  ∧ eval_expr get (ExtractElement ty ity v i) = do
      v <- get v; i <- get i;
      assert_type_eq ty (type_of_value v) "EXTRACTELEMENT array/vector";
      assert_type_eq ity (type_of_value i) "EXTRACTELEMENT index";
      off <- get_int_as_num "EXTRACTELEMENT index" i;
      C CONS [] <$> value_offset_get v off
    od
  ∧ eval_expr get (InsertElement ty ity v1 i v2) = do
      v1 <- get v1; i <- get i; v2 <- get v2;
      assert_type_eq ty (type_of_value v1) "INSERTELEMENT array/vector";
      assert_type_eq ity (type_of_value i) "INSERTELEMENT index";
      off <- get_int_as_num "INSERTELEMENT index" i;
      C CONS [] <$> value_offset_update v1 off v2
    od*)
    (* TODO: ShuffleVector *)
  ∧ eval_expr _ (ShuffleVector _ _ _ _ _) =
      K (Fail (NotImplemented "SHUFFLEVECTOR"))
  ∧ eval_expr ssa (GetIRef _ ref) = do
      rv <- HD (getc ssa ref);
      irv <- K (get_iref rv);
      return [irv]
    od
  ∧ eval_expr ssa (GetFieldIRef _ _ n ref) = do
      rv <- HD (getc ssa ref);
      rv' <- K (get_field_iref rv n);
      return [rv']
    od
  ∧ eval_expr ssa (GetElementIRef _ _ _ ref idx) = do
      rv <- HD (getc ssa ref);
      iv <- HD (getc ssa idx);
      rv' <- K (get_element_iref rv iv);
      return [rv']
    od
  ∧ eval_expr ssa (ShiftIRef _ _ _ ref off) = do
      rv <- HD (getc ssa ref);
      iv <- HD (getc ssa off);
      rv' <- K (shift_iref rv iv);
      return [rv']
    od
  ∧ eval_expr ssa (GetVarPartIRef _ _ ref) = do
      rv <- HD (getc ssa ref);
      ia <- K (get_iref_addr rv);
      case ia.ty of
      | Hybrid fts _ => C CONS [] <$> K (get_field_iref rv (LENGTH fts))
      | _ => K (Fail (InvalidIR "GETVARPARTIREF of non-hybrid"))
    od
`

(* μVM instructions are split up into operations before they are executed. *)
val compile_inst_def = Define`
    compile_inst (ssa : ssa_map) (Assign vs e) : opn list =
      [Nm (CURRY ZIP (LIST_BIND vs (MAP FST o ssa)) <$> eval_expr ssa e)]
  ∧ compile_inst (ssa : ssa_map) (Load dst _ _ src ord) : opn list =
      (let lds = MAP2 (λs d. Ld (s >>= K o get_iref_addr) (FST d))
                      (get ssa src) (ssa dst) in
       case ord of
       | NOT_ATOMIC => lds
       | RELAXED => lds
       | CONSUME => LIST_BIND lds (λld. [ld; Rec])
       | ACQUIRE => LIST_BIND lds (λld. [ld; Rec])
       | SEQ_CST => LIST_BIND lds (λld. [Com; Rec; ld; Rec])
       | _ => [Nm (K (Fail (InvalidIR "unsupported memory order for LOAD")))])
  ∧ compile_inst ssa (Store src _ _ dst ord) =
      (let sts = MAP2
         (λd s.  St ($, <$> (d >>= K o get_iref_addr) <*> s))
         (get ssa dst) (getc ssa src) in
       case ord of
       | NOT_ATOMIC => sts
       | RELAXED => sts
       | RELEASE => LIST_BIND sts (λst. [Com; st])
       | SEQ_CST => LIST_BIND sts (λst. [Com; st])
       | _ => [St (K (Fail (InvalidIR "unsupported memory order for STORE")))])
  ∧ compile_inst _ (CmpXchg _ _ _ _ _ _ _ _ _ _) =
      [Nm (K (Fail (NotImplemented "CMPXCHG")))]
  ∧ compile_inst _ (AtomicRMW _ _ _ _ _ _ _) =
      [Nm (K (Fail (NotImplemented "ATOMICRMW")))]
  ∧ compile_inst _ (Fence _) =
      [Nm (K (Fail (NotImplemented "FENCE")))]
  ∧ compile_inst ssa (New dst ty) =
      [Cmd (return (DoHeapAlloc ty NONE (FST (HD (ssa dst)))))]
  ∧ compile_inst ssa (Alloca dst ty) =
      [Cmd (return (DoStackAlloc ty NONE (FST (HD (ssa dst)))))]
  ∧ compile_inst ssa (NewHybrid dst ty _ len) =
      [Cmd ((λn. DoHeapAlloc ty (SOME n) (FST (HD (ssa dst))))
          <$> (HD (getc ssa len) >>= K o get_int_as_num "NEWHYBRID"))]
  ∧ compile_inst ssa (AllocaHybrid dst ty _ len) =
      [Cmd ((λn. DoStackAlloc ty (SOME n) (FST (HD (ssa dst))))
          <$> (HD (getc ssa len) >>= K o get_int_as_num "ALLOCAHYBRID"))]
  ∧ compile_inst ssa (NewThread dst stk tloc nsc) =
      [(* TODO: this*)]
`

val compile_terminst_def = Define`
    compile_terminst (lbl : block_label -> num) (ssa : ssa_map) (Ret vs) =
      [Cmd (DoPopStack o NOR <$> sequence (LIST_BIND vs (getc ssa)))]
  ∧ compile_terminst _ ssa (Throw v) =
      [Cmd (DoPopStack o EXC [] <$> HD (getc ssa v))]
  ∧ compile_terminst lbl ssa (Call dsts cdata rdata) =
     (let get_dst v =
        case v of
        | Var s => if MEM s dsts
                   then return (MAP (INL o FST) (ssa s))
                   else mapM (liftM INR) (get ssa s)
        | Const _ hd tl => return (MAP INR (hd::tl)) in
      let fname =
        case cdata.name of
        | INL n => do
            ref <- HD (get ssa n);
            (sig, name) <- K (get_funcref_data ref);
            (* TODO: Assert sig = cdata.signature *)
            return name
          od
        | INR name => return name in 
      let jmps = (λn e. SOME <|
          nor_jump := (lbl (FST rdata.normal_dest), n);
          exc_jump := OPTION_MAP (λed. lbl (FST ed), e) rdata.exceptional_dest
        |>) <$> (FLAT <$> mapM get_dst (SND rdata.normal_dest))
            <*> (FLAT <$> mapM get_dst (option_CASE rdata.exceptional_dest [] SND)) in
      [Cmd (DoPushStack <$> fname
                        <*> return cdata.signature
                        <*> sequence (LIST_BIND cdata.args (getc ssa))
                        <*> jmps)])
  ∧ compile_terminst _ ssa (TailCall cdata) =
     (let fname =
        case cdata.name of
        | INL n => do
            ref <- HD (get ssa n);
            (sig, name) <- K (get_funcref_data ref);
            (* TODO: Assert sig = cdata.signature *)
            return name
          od
        | INR name => return name in 
      [Cmd (DoPushStack <$> fname
                        <*> return cdata.signature
                        <*> sequence (LIST_BIND cdata.args (getc ssa))
                        <*> return NONE)])
  ∧ compile_terminst lbl ssa (Branch1 (b, vs)) =
      [Cmd (DoJump <$> return (lbl b) <*> sequence (LIST_BIND vs (getc ssa)))]
  ∧ compile_terminst lbl ssa (Branch2 p d1 d2) =
      [Cmd (HD (getc ssa p) >>= K o get_int1_as_bool >>= λis1.
         let (b, vs) = if is1 then d1 else d2 in
         DoJump <$> return (lbl b) <*> sequence (LIST_BIND vs (getc ssa)))]
  ∧ compile_terminst lbl ssa (Switch _ p default cases) =
      [Cmd (HD (getc ssa p) >>= λpv.
         let (b, vs) = option_CASE (ALOOKUP cases pv) default I in
         DoJump <$> return (lbl b) <*> sequence (LIST_BIND vs (getc ssa)))]
  ∧ compile_terminst _ _ (Watchpoint _ _ _) =
      [Cmd (K (Fail (NotImplemented "WATCHPOINT")))]
  ∧ compile_terminst _ _ (WPBranch _ _ _) =
      [Cmd (K (Fail (NotImplemented "WPBRANCH")))]
  ∧ compile_terminst _ _ (SwapStack _ _ _ _ _) =
      [Cmd (K (Fail (NotImplemented "SWAPSTACK")))]
  ∧ compile_terminst _ _ (CommInst _ _ _) =
      [Cmd (K (Fail (NotImplemented "COMMINST")))]
  ∧ compile_terminst _ _ (ExcClause _ _) =
      [Cmd (K (Fail (NotImplemented "EXC")))]
`

val compile_block_def = Define`
  compile_block (lbl : block_label -> num) (block : bblock) : opn_block =
    let next n = (n, SUC n) in
    let alloc_cells (cs, ssas, n) (ssa, ty) =
      let (cs', n') = mem_cell_types NONE ty next n in
      (cs ++ MAP SND cs', ssas |+ (ssa, MAP (I ## FST) cs'), n') in
    let (cs, ssas, _) = FOLDL alloc_cells ([], FEMPTY, 0) (block_types block) in
    <|
      params := TAKE (SUM (MAP (LENGTH o FAPPLY ssas o FST) block.args)) cs;
      has_exc_param := IS_SOME block.exc;
      opns :=
        LIST_BIND block.body (compile_inst (FAPPLY ssas)) ++
        compile_terminst lbl (FAPPLY ssas) block.exit
    |>
`

val compile_function_def = Define`
  compile_function (func : function) : opn_block spt =
    let lbl_alist =
      MAP (λn. (FST (EL n func.blocks), n))
          (GENLIST I (LENGTH func.blocks)) in
    fromList (MAP (compile_block (THE o ALOOKUP lbl_alist) o SND) func.blocks)
`

val compile_program_def = Define`
  compile_program (env : environment) : program =
    compile_function o_f
    env.functions f_o_f
    THE o_f
    RRESTRICT env.func_versions IS_SOME
`

val _ = export_theory()