open import Pervasives_extra
open import Sail_instr_kinds
open import Sail_values

(* 'a is result type *)

type memstate = map integer memory_byte
type tagstate = map integer bitU
(* type regstate = map string (vector bitU) *)

type sequential_state 'regs =
  <| regstate : 'regs;
     memstate : memstate;
     tagstate : tagstate;
     write_ea : maybe (write_kind * integer * integer);
     last_exclusive_operation_was_load : bool|>

val init_state : forall 'regs. 'regs -> sequential_state 'regs
let init_state regs =
  <| regstate = regs;
     memstate = Map.empty;
     tagstate = Map.empty;
     write_ea = Nothing;
     last_exclusive_operation_was_load = false |>

type ex 'e =
  | Exit
  | Failure of string
  | Throw of 'e

type result 'a 'e =
  | Value of 'a
  | Ex of (ex 'e)

(* State, nondeterminism and exception monad with result value type 'a
   and exception type 'e. *)
type monadS 'regs 'a 'e = sequential_state 'regs -> list (result 'a 'e * sequential_state 'regs)

val returnS : forall 'regs 'a 'e. 'a -> monadS 'regs 'a 'e
let returnS a s = [(Value a,s)]

val bindS : forall 'regs 'a 'b 'e. monadS 'regs 'a 'e -> ('a -> monadS 'regs 'b 'e) -> monadS 'regs 'b 'e
let bindS m f (s : sequential_state 'regs) =
  List.concatMap (function
                  | (Value a, s') -> f a s'
                  | (Ex e, s') -> [(Ex e, s')]
                  end) (m s)

val seqS: forall 'regs 'b 'e. monadS 'regs unit 'e -> monadS 'regs 'b 'e -> monadS 'regs 'b 'e
let seqS m n = bindS m (fun (_ : unit) -> n)

val exitS : forall 'regs 'e 'a. unit -> monadS 'regs 'a 'e
let exitS () s = [(Ex Exit, s)]

val failS : forall 'regs 'a 'e. string -> monadS 'regs 'a 'e
let failS msg s = [(Ex (Failure msg), s)]

val throwS : forall 'regs 'a 'e. 'e -> monadS 'regs 'a 'e
let throwS e s = [(Ex (Throw e), s)]

val try_catchS : forall 'regs 'a 'e1 'e2. monadS 'regs 'a 'e1 -> ('e1 -> monadS 'regs 'a 'e2) ->  monadS 'regs 'a 'e2
let try_catchS m h s =
  List.concatMap (function
                  | (Value a, s') -> returnS a s'
                  | (Ex (Throw e), s') -> h e s'
                  | (Ex Exit, s') -> [(Ex Exit, s')]
                  | (Ex (Failure msg), s') -> [(Ex (Failure msg), s')]
                  end) (m s)

val assert_expS : forall 'regs 'e. bool -> string -> monadS 'regs unit 'e
let assert_expS exp msg = if exp then returnS () else failS msg

(* For early return, we abuse exceptions by throwing and catching
   the return value. The exception type is "either 'r 'e", where "Right e"
   represents a proper exception and "Left r" an early return of value "r". *)
type monadSR 'regs 'a 'r 'e = monadS 'regs 'a (either 'r 'e)

val early_returnS : forall 'regs 'a 'r 'e. 'r -> monadSR 'regs 'a 'r 'e
let early_returnS r = throwS (Left r)

val catch_early_returnS : forall 'regs 'a 'e. monadSR 'regs 'a 'a 'e -> monadS 'regs 'a 'e
let catch_early_returnS m =
  try_catchS m
    (function
      | Left a -> returnS a
      | Right e -> throwS e
     end)

(* Lift to monad with early return by wrapping exceptions *)
val liftSR : forall 'a 'r 'regs 'e. monadS 'regs 'a 'e -> monadSR 'regs 'a 'r 'e
let liftSR m = try_catchS m (fun e -> throwS (Right e))

(* Catch exceptions in the presence of early returns *)
val try_catchSR : forall 'regs 'a 'r 'e1 'e2. monadSR 'regs 'a 'r 'e1 -> ('e1 -> monadSR 'regs 'a 'r 'e2) ->  monadSR 'regs 'a 'r 'e2
let try_catchSR m h =
  try_catchS m
    (function
      | Left r -> throwS (Left r)
      | Right e -> h e
     end)

val range : integer -> integer -> list integer
let rec range i j =
  if j < i then []
  else if i = j then [i]
  else i :: range (i+1) j

val get_regS : forall 'regs 'rv 'a. sequential_state 'regs -> register_ref 'regs 'rv 'a -> 'a
let get_regS state reg = reg.read_from state.regstate

val set_regS : forall 'regs 'rv 'a. sequential_state 'regs -> register_ref 'regs 'rv 'a -> 'a -> sequential_state 'regs
let set_regS state reg v =
  <| state with regstate = reg.write_to state.regstate v |>


val read_memS : forall 'regs 'e. read_kind -> integer -> integer -> monadS 'regs (list memory_byte) 'e
let read_memS read_kind addr sz s =
  (*let addr = unsigned (bitv_of_address_lifted addr) in
  let sz = integerFromNat sz in*)
  let addrs = range addr (addr+sz-1) in
  match just_list (List.map (fun addr -> Map.lookup addr s.memstate) addrs) with
    | Just mem_val ->
       let s' =
         if read_is_exclusive read_kind
         then <| s with last_exclusive_operation_was_load = true |>
         else s
       in
       returnS (List.reverse mem_val) s'
    | Nothing -> failS "read_memS" s
  end

(* caps are aligned at 32 bytes *)
let cap_alignment = (32 : integer)

val read_tagS : forall 'regs 'a 'e. Bitvector 'a => read_kind -> 'a -> monadS 'regs bitU 'e
let read_tagS read_kind addr state =
  let addr = (unsigned addr) / cap_alignment in
  let tag = match (Map.lookup addr state.tagstate) with
    | Just t -> t
    | Nothing -> B0
  end in
  if read_is_exclusive read_kind
  then [(Value tag, <| state with last_exclusive_operation_was_load = true |>)]
  else [(Value tag, state)]

val excl_resultS : forall 'regs 'e. unit -> monadS 'regs bool 'e
let excl_resultS () state =
  let success =
    (Value true,  <| state with last_exclusive_operation_was_load = false |>) in
  (Value false, state) :: if state.last_exclusive_operation_was_load then [success] else []

val write_mem_eaS : forall 'regs 'e. write_kind -> integer -> integer -> monadS 'regs unit 'e
let write_mem_eaS write_kind addr sz state =
  (*let addr = unsigned (bitv_of_address_lifted addr) in
  let sz = integerFromNat sz in*)
  [(Value (), <| state with write_ea = Just (write_kind, addr, sz) |>)]

val write_mem_valS : forall 'regs 'e. list memory_byte -> monadS 'regs bool 'e
let write_mem_valS v state =
  let (_,addr,sz) = match state.write_ea with
    | Nothing -> failwith "write ea has not been announced yet"
    | Just write_ea -> write_ea end in
  let addrs = range addr (addr+sz-1) in
  (*let v = external_mem_value (bits_of v) in*)
  let addresses_with_value = List.zip addrs (List.reverse v) in
  let memstate = List.foldl (fun mem (addr,v) -> Map.insert addr v mem)
                            state.memstate addresses_with_value in
  [(Value true,  <| state with memstate = memstate |>)]

val write_tagS : forall 'regs 'e. bitU -> monadS 'regs bool 'e
let write_tagS t state =
  let (_,addr,_) = match state.write_ea with
    | Nothing -> failwith "write ea has not been announced yet"
    | Just write_ea -> write_ea end in
  let taddr = addr / cap_alignment in
  let tagstate = Map.insert taddr t state.tagstate in
  [(Value true,  <| state with tagstate = tagstate |>)]

val read_regS : forall 'regs 'rv 'a 'e. register_ref 'regs 'rv 'a -> monadS 'regs 'a 'e
let read_regS reg s = [(Value (reg.read_from s.regstate), s)]

(* TODO
let read_reg_range reg i j state =
  let v = slice (get_reg state (name_of_reg reg)) i j in
  [(Value (vec_to_bvec v),state)]
let read_reg_bit reg i state =
  let v = access (get_reg state (name_of_reg reg)) i in
  [(Value v,state)]
let read_reg_field reg regfield =
  let (i,j) = register_field_indices reg regfield in
  read_reg_range reg i j
let read_reg_bitfield reg regfield =
  let (i,_) = register_field_indices reg regfield in
  read_reg_bit reg i *)

val read_regvalS : forall 'regs 'rv 'e.
  register_accessors 'regs 'rv -> string -> monadS 'regs 'rv 'e
let read_regvalS (read, _) reg s =
  match read reg s.regstate with
    | Just v ->  returnS v s
    | Nothing -> failS ("read_regvalS " ^ reg) s
  end

val write_regvalS : forall 'regs 'rv 'e.
  register_accessors 'regs 'rv -> string -> 'rv -> monadS 'regs unit 'e
let write_regvalS (_, write) reg v s =
  match write reg v s.regstate with
    | Just rs' -> returnS () (<| s with regstate = rs' |>)
    | Nothing ->  failS ("write_regvalS " ^ reg) s
  end

val write_regS : forall 'regs 'rv 'a 'e. register_ref 'regs 'rv 'a -> 'a -> monadS 'regs unit 'e
let write_regS reg v state =
  [(Value (), <| state with regstate = reg.write_to state.regstate v |>)]

(* TODO
val update_reg : forall 'regs 'rv 'a 'b 'e. register_ref 'regs 'rv 'a -> ('a -> 'b -> 'a) -> 'b -> monadS 'regs unit 'e
let update_reg reg f v state =
  let current_value = get_reg state reg in
  let new_value = f current_value v in
  [(Value (), set_reg state reg new_value)]

let write_reg_field reg regfield = update_reg reg regfield.set_field

val update_reg_range : forall 'regs 'rv 'a 'b. Bitvector 'a, Bitvector 'b => register_ref 'regs 'rv 'a -> integer -> integer -> 'a -> 'b -> 'a
let update_reg_range reg i j reg_val new_val = set_bits (reg.is_inc) reg_val i j (bits_of new_val)
let write_reg_range reg i j = update_reg reg (update_reg_range reg i j)

let update_reg_pos reg i reg_val x = update_list reg.is_inc reg_val i x
let write_reg_pos reg i = update_reg reg (update_reg_pos reg i)

let update_reg_bit reg i reg_val bit = set_bit (reg.is_inc) reg_val i (to_bitU bit)
let write_reg_bit reg i = update_reg reg (update_reg_bit reg i)

let update_reg_field_range regfield i j reg_val new_val =
  let current_field_value = regfield.get_field reg_val in
  let new_field_value = set_bits (regfield.field_is_inc) current_field_value i j (bits_of new_val) in
  regfield.set_field reg_val new_field_value
let write_reg_field_range reg regfield i j = update_reg reg (update_reg_field_range regfield i j)

let update_reg_field_pos regfield i reg_val x =
  let current_field_value = regfield.get_field reg_val in
  let new_field_value = update_list regfield.field_is_inc current_field_value i x in
  regfield.set_field reg_val new_field_value
let write_reg_field_pos reg regfield i = update_reg reg (update_reg_field_pos regfield i)

let update_reg_field_bit regfield i reg_val bit =
  let current_field_value = regfield.get_field reg_val in
  let new_field_value = set_bit (regfield.field_is_inc) current_field_value i (to_bitU bit) in
  regfield.set_field reg_val new_field_value
let write_reg_field_bit reg regfield i = update_reg reg (update_reg_field_bit regfield i)*)