open import Pervasives_extra
open import Vector
open import Arch

(* 'a is result type, 'e is error type *)
type State 's 'e 'a = 's -> (either 'a 'e * 's)

type memstate = map integer (list bit)

type state =
  <| regstate : regstate;
     memstate : memstate;
     writeEA : integer * integer|>

let compose f g x = f (g x)

val ints_aux : list integer -> integer -> list integer
let rec ints_aux acc x = 
  if x = 0 then []
  else ints_aux ((x - 1) :: acc) (x - 1)

let ints = ints_aux []

val return : forall 's 'e 'a. 'a -> State 's 'e 'a
let return a s = (Left a,s)

val bind : forall 's 'e 'a 'b. State 's 'e 'a -> ('a -> State 's 'e 'b) -> State 's 'e 'b
let bind m f s = match m s with
  | (Left a,s') -> f a s'
  | (Right error,s') -> (Right error,s')
  end

val exit : forall 's 'e. 'e -> State 's 'e unit
let exit e s = (Right e,s)

let (>>=) = bind

val (>>): forall 's 'e 'a 'b. State 's 'e 'a -> State 's 'e 'b -> State 's 'e 'b
let (>>) m n = m >>= fun _ -> n

val read_writeEA : forall 'e. unit -> State state 'e (integer * integer)
let read_writeEA () s = (Left s.writeEA,s)

val write_writeEA : forall 'e. integer -> integer -> State state 'e unit
let write_writeEA addr l s = (Left (), <| s with writeEA = (addr,l) |>)

val byte_chunks : forall 'a. integer -> list 'a -> list (list 'a)
let rec byte_chunks n l =
  if n = 0 then []
  else
    match l with
    | a::b::c::d::e::f::g::h::rest -> [a;b;c;d;e;f;g;h] :: byte_chunks (n - 1) rest
    | _ -> failwith "byte_chunks not given enough bits"
    end

val read_regstate : state -> register -> vector bit
let read_regstate s = read_regstate_aux s.regstate

val write_regstate : state -> register -> vector bit -> state
let write_regstate s reg v = <| s with regstate = (write_regstate_aux s.regstate reg v) |>

val read_memstate : memstate -> integer -> list bit
let read_memstate s addr = findWithDefault addr (replicate 8 Undef) s

val write_memstate : memstate -> (integer * list bit) -> memstate
let write_memstate s (addr,bits) = Map.insert addr bits s

val read_memory : forall 'e. integer -> integer -> State state 'e (vector bit)
let read_memory addr l s =
  let bytes = List.map (compose (read_memstate s.memstate) ((+) addr)) (ints l) in
  (Left (V (foldl (++) [] bytes) 0 defaultDir),s)

val write_memory : forall 'e. integer -> integer -> vector bit -> State state 'e unit
let write_memory addr l (V bits _ _) s =
  let addrs = List.map ((+) addr) (ints l) in
  let bytes = byte_chunks l bits in
  (Left (), <| s with memstate = foldl write_memstate s.memstate (zip addrs bytes) |>)

val read_reg_range : forall 'e. register -> integer -> integer (*(nat * nat)*) -> State state 'e (vector bit)
let read_reg_range reg i j s =
  let v = slice (read_regstate s reg) i j in
  (Left v,s)

val read_reg_bit : forall 'e. register -> integer (*nat*) -> State state 'e bit
let read_reg_bit reg i s =
  let v = access (read_regstate s reg) i in
  (Left v,s)

val write_reg_range : forall 'e. register -> integer -> integer (*(nat * nat)*) -> vector bit -> State state 'e unit
let write_reg_range reg i j v s =
  let v' = update (read_regstate s reg) i j v in
  let s' = write_regstate s reg v' in
  (Left (),s')

val write_reg_bit : forall 'e. register -> integer (*nat*) -> bit -> State state 'e unit
let write_reg_bit reg i bit s =
  let v = read_regstate s reg in
  let v' = update_pos v i bit in
  let s' = write_regstate s reg v' in
  (Left (),s')

val read_reg : forall 'e. register -> State state 'e (vector bit)
let read_reg reg s =
  let v = read_regstate s reg in
  (Left v,s)

val write_reg : forall 'e. register -> vector bit -> State state 'e unit
let write_reg reg v s =
  let s' = write_regstate s reg v in
  (Left (),s')
    

val foreach_inc :  forall 's 'e 'vars. (integer * integer * integer) (*(nat * nat * nat)*) -> 'vars ->
                  (integer (*nat*) -> 'vars -> (unit * 'vars)) -> (unit * 'vars)
let rec foreach_inc (i,stop,by) vars body = 
  if i <= stop
  then
    let (_,vars) = body i vars in
    foreach_inc (i + by,stop,by) vars body
  else ((),vars)


val foreach_dec : forall 's 'e 'vars. (integer * integer * integer) (*(nat * nat * nat)*) -> 'vars ->
                  (integer (*nat*) -> 'vars -> (unit * 'vars)) -> (unit * 'vars)
let rec foreach_dec (i,stop,by) vars body = 
  if i >= stop
  then
    let (_,vars) = body i vars in
    foreach_dec (i - by,stop,by) vars body
  else ((),vars)


val foreachM_inc : forall 's 'e 'vars. (integer * integer * integer) -> 'vars ->
                  (integer -> 'vars -> State 's 'e (unit * 'vars)) -> State 's 'e (unit * 'vars)
let rec foreachM_inc (i,stop,by) vars body = 
  if i <= stop
  then
    body i vars >>= fun (_,vars) ->
    foreachM_inc (i + by,stop,by) vars body
  else return ((),vars)


val foreachM_dec : forall 's 'e 'vars. (integer * integer * integer) -> 'vars ->
                  (integer -> 'vars -> State 's 'e (unit * 'vars)) -> State 's 'e (unit * 'vars)
let rec foreachM_dec (i,stop,by) vars body = 
  if i >= stop
  then
    body i vars >>= fun (_,vars) ->
    foreachM_dec (i - by,stop,by) vars body
  else return ((),vars)

val read_reg_field : forall 'e. register -> register_field -> State state 'e (vector bit)
let read_reg_field reg rfield = uncurry (read_reg_range reg) (field_indices rfield)

val write_reg_field : forall 'e. register -> register_field -> vector bit -> State state 'e unit
let write_reg_field reg rfield = uncurry (write_reg_range reg) (field_indices rfield)

val read_reg_bitfield : forall 'e. register -> register_bitfield -> State state 'e bit
let read_reg_bitfield reg rbit = read_reg_bit reg (field_index_bit rbit)

val write_reg_bitfield : forall 'e. register -> register_bitfield -> bit -> State state 'e unit
let write_reg_bitfield reg rbit = write_reg_bit reg (field_index_bit rbit)

val write_reg_field_range : forall 'e. register -> register_field -> integer -> integer ->
                            vector bit -> State state 'e unit
let write_reg_field_range reg rfield i j v =
  let (i',j') = field_indices rfield in
  write_reg_range reg i' j' v

val write_reg_field_bit : forall 'e. register -> register_field -> integer ->
                            bit -> State state 'e unit
let write_reg_field_bit reg rfield i v =
  let (i',_) = field_indices rfield in
  write_reg_bit reg (i + i') v



let length l = integerFromNat (length l)

let write_two_regs r1 r2 vec =
  let size = length_reg r1 in
  let start = get_start vec in
  let vsize = length vec in
  let r1_v = slice vec start ((if defaultDir then size - start else start - size) - 1) in
  let r2_v =
    (slice vec)
      (if defaultDir then size - start else start - size)
      (if defaultDir then vsize - start else start - vsize) in
  write_reg r1 r1_v >> write_reg r2 r2_v

let read_two_regs r1 r2 =
  read_reg r1 >>= fun v1 ->
  read_reg r2 >>= fun v2 ->
  return (v1 ^^ v2)

type M 'e 'a = State state 'e 'a