open import Pervasives
open import Interp
open import Interp_ast
import Assert_extra Maybe_extra (* For failwith for error reporting while debugging; and for fromJust when we know it's not Nothing *)
open import Num
open import List
open import Word
open import Bool

let ignore_sail x = V_lit (L_aux L_unit Unknown) ;;

let compose f g x = f (V_tuple [g x]) ;;

let zeroi = integerFromNat 0
let onei = integerFromNat 1
let twoi = integerFromNat 2

let rec sparse_walker update ni processed_length length ls df =
  if processed_length = length 
  then []
  else match ls with
  | [] -> replicate (natFromInteger (length - processed_length)) df
  | (i,v)::ls -> 
    if ni = i 
    then v::(sparse_walker update (update ni) (processed_length + 1) length ls df)
    else df::(sparse_walker update (update ni) (processed_length + 1) length ((i,v)::ls) df)
end

let rec fill_in_sparse v = match v with
  | V_vector_sparse first length inc ls df -> 
    V_vector first inc (sparse_walker (if inc then (fun (x: integer) -> x + onei) else (fun (x: integer) -> x - onei)) first zeroi length ls df)
  | V_unknown -> V_unknown
  | V_track v r -> taint (fill_in_sparse v) r
end    

let is_one v = match v with 
  | V_lit (L_aux b lb) -> V_lit (L_aux (if b = L_one then L_true else L_false) lb)
  | V_track (V_lit (L_aux b lb)) r -> V_track (V_lit (L_aux (if b = L_one then L_true else L_false) lb)) r
  | V_track V_unkown _ -> v
  | V_unknown -> v
end ;;

let rec lt_range (V_tuple[v1;v2]) = match (v1,v2) with 
  | (V_lit (L_aux (L_num l1) lr),V_lit (L_aux (L_num l2) ll)) ->
    if l1 < l2
    then V_lit (L_aux L_one Unknown)
    else V_lit (L_aux L_zero Unknown)
  | (V_track v1 r1, V_track v2 r2) -> 
    taint (lt_range (V_tuple [v1;v2])) (r1++r2)
  | (V_track v1 r1, v2) ->
    taint (lt_range (V_tuple [v1;v2])) r1
  | (v1,V_track v2 r2) ->
    taint (lt_range (V_tuple [v1;v2])) r2
  | (V_unknown,_) -> V_unknown
  | (_,V_unknown) -> V_unknown
end ;;

let bit_to_bool b = match b with
  | V_lit (L_aux L_zero _) -> false
  | V_track (V_lit (L_aux L_zero _)) _ -> false
  | V_lit (L_aux L_one _) -> true
  | V_track (V_lit (L_aux L_one _)) _ -> true
(* PS HACK TO MAKE BUILD *)
  | _ -> false  (* ARGH - SHOULDN'T HAVE TO CHOOSE *)
(* END HACK *)
  end ;;
let bool_to_bit b = match b with
    false -> V_lit (L_aux L_zero Unknown)
  | true -> V_lit (L_aux L_one Unknown)
  end ;;

let bitwise_not_bit v = 
  let lit_not (L_aux l loc) = match l with
    | L_zero -> (V_lit (L_aux L_one loc))
    | L_one  -> (V_lit (L_aux L_zero loc)) end in
  match v with
  | V_lit lit -> lit_not lit
  | V_track (V_lit lit) r -> V_track (lit_not lit) r 
  | V_unknown -> v
  | V_track V_unknown r -> v
end;; 

let rec bitwise_not v =
  match v with
  | V_vector idx inc v ->
    V_vector idx inc (List.map bitwise_not_bit v)
  | V_unknown -> v
  | V_track bv r -> taint (bitwise_not bv) r
end ;;

let rec bitwise_binop_bit op (V_tuple [x; y]) = match (x,y) with
  | (V_track x rx,V_track y ry) ->
    taint ((bitwise_binop_bit op) (V_tuple[x; y])) (rx ++ ry)
  | (V_track x rx,y) ->
    taint ((bitwise_binop_bit op) (V_tuple[x;y])) rx
  | (x,V_track y ry) ->
    taint ((bitwise_binop_bit op) (V_tuple[x;y])) ry
  | (V_unknown,_) -> V_unknown
  | (_,V_unknown) -> V_unknown
  | _ -> bool_to_bit (op (bit_to_bool x) (bit_to_bool y))
end ;;

let rec bitwise_binop op (V_tuple [v1;v2]) =
  match (v1,v2) with
  | (V_vector idx inc v, V_vector idx' inc' v') ->
  (* typechecker ensures inc = inc', idx = idx' and length v = length v' *)
    V_vector idx inc (List.map (fun (x,y) -> (bitwise_binop_bit op (V_tuple[x; y]))) (List.zip v v'))
  | (V_track v1 r1, V_track v2 r2) ->
    taint (bitwise_binop op (V_tuple [v1;v2])) (r1 ++ r2)
  | (V_track v1 r1,v2) ->
    taint (bitwise_binop op (V_tuple [v1;v2])) r1
  | (v1,V_track v2 r2) ->
    taint (bitwise_binop op (V_tuple [v1;v2])) r2
  | (V_unknown,_) -> V_unknown
  | (_,V_unknown) -> V_unknown
end ;;

(* BitSeq expects LSB first.
 * By convention, MSB is on the left, so increasing = Big-Endian (MSB0),
 * hence MSB first.
 * http://en.wikipedia.org/wiki/Bit_numbering *)
let rec to_num signed v = 
  match v with
  | (V_vector idx inc l) ->
  (* Word library in Lem expects bitseq with LSB first *)
(*TODO: Kathy think more
  We thought a reverse was needed due to above comments only in inc case.
  However, we seem to be interpresting bit vectors such that reverse is always needed
  and despite numbering MSB is on the left.
*)
  let l = (*if inc then reverse l else l*) reverse l in
  (* Make sure the last bit is a zero to force unsigned numbers *)
  let l = if signed then l else l ++ [V_lit (L_aux L_zero Unknown)] in
  V_lit(L_aux (L_num(integerFromBitSeq (Maybe_extra.fromJust (bitSeqFromBoolList (map bit_to_bool l))))) Unknown)
  | V_unknown -> V_unknown
  | V_track v r -> taint (to_num signed v) r
end ;;

let rec to_vec_inc (V_tuple[v1;v2]) = match (v1,v2) with
  | (V_lit(L_aux (L_num len) _), (V_lit(L_aux (L_num n) ln))) ->
    let l = boolListFrombitSeq (natFromInteger len) (bitSeqFromInteger Nothing n) in
    V_vector 0 true (map bool_to_bit (reverse l)) 
  | (V_track v1 r1,V_track v2 r2) -> taint (to_vec_inc (V_tuple[v1;v2])) (r1++r2)
  | (V_track v1 r1,v2) -> taint (to_vec_inc (V_tuple[v1;v2])) r1
  | (v1,V_track v2 r2) -> taint (to_vec_inc (V_tuple[v1;v2])) r2
  | (V_unknown,_) -> V_unknown
  | (_,V_unknown) -> V_unknown
end
;;
let rec to_vec_dec (V_tuple([v1;v2])) = match (v1,v2) with 
  | (V_lit(L_aux (L_num len) _), (V_lit(L_aux (L_num n) ln))) ->
    let l = boolListFrombitSeq (natFromInteger len) (bitSeqFromInteger Nothing n) in
    V_vector (len - 1) false (map bool_to_bit (reverse l)) 
  | (V_track v1 r1,V_track v2 r2) -> taint (to_vec_dec (V_tuple[v1;v2])) (r1++r2)
  | (V_track v1 r1,v2) -> taint (to_vec_dec (V_tuple[v1;v2])) r1
  | (v1,V_track v2 r2) -> taint (to_vec_dec (V_tuple[v1;v2])) r2
  | (V_unknown,_) -> V_unknown
  | (_,V_unknown) -> V_unknown
end
;;
let to_vec ord len n =
  if ord
  then to_vec_inc (V_tuple ([V_lit(L_aux (L_num (integerFromNat len)) Interp_ast.Unknown); n]))
  else to_vec_dec (V_tuple ([V_lit(L_aux (L_num (integerFromNat len)) Interp_ast.Unknown); n]))
;;

let rec exts (V_tuple[v1;v]) = match v1 with
  | V_lit(L_aux (L_num len) _) -> (match v with
      | V_vector _ inc _  -> to_vec inc (natFromInteger len) (to_num true v)
      | V_unknown -> V_unknown
      | V_track v r2 -> taint (exts (V_tuple[v1;v])) r2 end)
  | V_unknown -> v1
  | V_track v1 r1 -> taint (exts (V_tuple[v1;v])) r1
end
;;

let eq (V_tuple [x; y]) = V_lit (L_aux (if value_eq x y then L_one else L_zero) Unknown) ;;
(* XXX interpret vectors as unsigned numbers for equality *)
let eq_vec_range (V_tuple [v; r]) = eq (V_tuple [to_num false v; r]) ;;
let eq_range_vec (V_tuple [r; v]) = eq (V_tuple [r; to_num false v]) ;;

let rec neg v = match v with
  | V_lit (L_aux arg la) -> 
    V_lit (L_aux (match arg with
      | L_one -> L_zero
      | L_zero -> L_one end) la) 
  | V_unknown -> V_unknown
  | V_track v r -> taint (neg v) r
  | V_tuple [v] -> neg v
end 
;;

let neq = compose neg eq ;;

let rec arith_op op (V_tuple args) = match args with
  | [V_lit(L_aux (L_num x) lx); V_lit(L_aux (L_num y) ly)] -> V_lit(L_aux (L_num (op x y)) lx)
  | [V_track v1 r1;V_track v2 r2] -> taint (arith_op op (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1;v2] -> taint (arith_op op (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] -> taint (arith_op op (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  end ;;
let rec arith_op_vec op size (V_tuple args) = match args with
  | [(V_vector b ord cs as l1);(V_vector _ _ _ as l2)] ->
    let (l1',l2') = (to_num false l1,to_num false l2) in
    let n = arith_op op (V_tuple [l1';l2']) in
    to_vec ord ((List.length cs) * size) n
  | [V_track v1 r1;V_track v2 r2] ->
    taint (arith_op_vec op size (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1;v2] ->
    taint (arith_op_vec op size (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] ->
    taint (arith_op_vec op size (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
end ;;
let rec arith_op_range_vec op size (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (arith_op_range_vec op size (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] ->
    taint (arith_op_range_vec op size (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] -> 
    taint (arith_op_range_vec op size (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [n; (V_vector _ ord cs as l2)] ->
    arith_op_vec op size (V_tuple [(to_vec ord (List.length cs) n);l2])
end ;;
let rec arith_op_vec_range op size (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (arith_op_vec_range op size (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] ->
    taint (arith_op_vec_range op size (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] -> 
    taint (arith_op_vec_range op size (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [(V_vector _ ord cs as l1);n] ->
    arith_op_vec op size (V_tuple [l1;(to_vec ord (List.length cs) n)])
end ;;
let rec arith_op_range_vec_range op (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (arith_op_range_vec_range op (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] ->
    taint (arith_op_range_vec_range op (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] -> 
    taint (arith_op_range_vec_range op (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [n;(V_vector _ ord _ as l2)] ->
    arith_op op (V_tuple [n;(to_num false l2)])
end ;;
let rec arith_op_vec_range_range op (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (arith_op_vec_range_range op (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] ->
    taint (arith_op_vec_range_range op (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] -> 
    taint (arith_op_vec_range_range op (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [(V_vector _ ord _ as l1);n] ->
    arith_op op (V_tuple [(to_num false l1);n])
end ;;

let rec compare_op op (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (compare_op op (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1;v2] ->
    taint (compare_op op (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] ->
    taint (compare_op op (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [V_lit(L_aux (L_num x) lx); V_lit(L_aux (L_num y) ly)] -> 
    if (op x y)
    then V_lit(L_aux L_one lx)
    else V_lit(L_aux L_zero lx)
end ;;
let rec compare_op_vec op (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] ->
    taint (compare_op_vec op (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1;v2] ->
    taint (compare_op_vec op (V_tuple [v1;v2])) r1
  | [v1;V_track v2 r2] ->
    taint (compare_op_vec op (V_tuple [v1;v2])) r2
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [((V_vector _ _ _) as l1);((V_vector _ _ _) as l2)] ->
    let (l1',l2') = (to_num true l1, to_num true l2) in
    compare_op op (V_tuple[l1';l2'])
end ;;

let rec duplicate (V_tuple args) = match args with
  | [V_track v1 r1;V_track v2 r2] -> taint (duplicate (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] -> taint (duplicate (V_tuple[v1;v2])) r1 
  | [v1;V_track v2 r2] -> taint (duplicate (V_tuple[v1;v2])) r2 
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  | [(V_lit _ as v);(V_lit (L_aux (L_num n) _))] ->
    (V_vector 0 true (List.replicate (natFromInteger n) v))
end

let rec vec_concat (V_tuple args) = match args with
  | [V_vector n d l; V_vector n' d' l'] ->
      (* XXX d = d' ? dropping n' ? *)
      V_vector n d (l ++ l')
  | [V_lit l; (V_vector n d l' as x)] -> vec_concat (V_tuple [litV_to_vec l d; x])
  | [(V_vector n d l' as x); V_lit l] -> vec_concat (V_tuple [x; litV_to_vec l d]) 
  | [V_track v1 r1;V_track v2 r2] -> taint (vec_concat (V_tuple [v1;v2])) (r1++r2)
  | [V_track v1 r1; v2] -> taint (vec_concat (V_tuple[v1;v2])) r1 
  | [v1; V_track v2 r2] -> taint (vec_concat (V_tuple[v1;v2])) r2 
  | [V_unknown;_] -> V_unknown
  | [_;V_unknown] -> V_unknown
  
  end ;;

let rec v_length v = match v with
  | V_vector n d l -> V_lit (L_aux (L_num (integerFromNat (List.length l))) Unknown)
  | V_unknown -> V_unknown
  | V_track v r -> taint (v_length v) r
end ;;

let function_map = [
  ("ignore", ignore_sail);            
  ("length", v_length);
  ("add", arith_op (+));
  ("add_vec", arith_op_vec (+) 1);
  ("add_vec_range", arith_op_vec_range (+) 1);
  ("add_vec_range_range", arith_op_vec_range_range (+));
  ("add_range_vec", arith_op_range_vec (+) 1);
  ("add_range_vec_range", arith_op_range_vec_range (+));
  ("minus", arith_op (-));
  ("minus_vec", arith_op_vec (-) 1);
  ("minus_vec_range", arith_op_vec_range (-) 1);
  ("minus_range_vec", arith_op_range_vec (-) 1);
  ("minus_vec_range_range", arith_op_vec_range_range (-));
  ("minus_range_vec_range", arith_op_range_vec_range (-));
  ("multiply", arith_op ( * ));
  ("multiply_vec", arith_op_vec ( * ) 2);
  ("mult_range_vec", arith_op_range_vec ( * ) 2);
  ("mult_vec_range", arith_op_vec_range ( * ) 2);
  ("mod", arith_op (mod));
  ("mod_vec", arith_op_vec (mod) 1);
  ("mod_vec_range", arith_op_vec_range (mod) 1);
  ("eq", eq);
  ("eq_vec_range", eq_vec_range);
  ("eq_range_vec", eq_range_vec);
  ("neq", neq);
  ("vec_concat", vec_concat);
  ("is_one", is_one);
  ("to_num_inc", to_num false);
  ("to_num_dec", to_num false);
  ("EXTS", exts);
  ("to_vec_inc", to_vec_inc);
  ("to_vec_dec", to_vec_dec);
  ("bitwise_not", bitwise_not);
  ("bitwise_not_bit", bitwise_not_bit);
  ("bitwise_and", bitwise_binop (&&));
  ("bitwise_or", bitwise_binop (||));
  ("bitwise_xor", bitwise_binop xor);
  ("bitwise_and_bit", bitwise_binop_bit (&&));
  ("bitwise_or_bit", bitwise_binop_bit (||));
  ("bitwise_xor_bit", bitwise_binop_bit xor);
  ("lt", compare_op (<));
  ("gt", compare_op (>));
  ("lt_vec", compare_op_vec (<));
  ("gt_vec", compare_op_vec (>));
  ("duplicate", duplicate);
] ;;

let eval_external name v = match List.lookup name function_map with
  | Just f -> f v
  | Nothing -> Assert_extra.failwith ("missing library function " ^ name)
  end