array.ml

type sort = Bool | Int | Array of sort 

type op = Plus | Times | Minus 

type var = string

type 'a termF = 
  | Var of var
  | Lit of int
  | Op of op * 'a * 'a 
  | Select of 'a * 'a          (* a[i] *)
  | Store of 'a * 'a * 'a    (* a[i] = t *)

let map_termF f (x : 'a termF) = 
  match x with
  | Var x -> Var x
  | Lit n -> Lit n 
  | Op (op, a1, a2) -> Op(op, f a1, f a2)
  | Select (a1, a2) -> Select(f a1, f a2)
  | Store (a1, a2, a3) -> Store(f a1, f a2, f a3)

type term = In of term termF

module TermBuild = struct
  let var x = In(Var x)
  let lit n = In(Lit n)
  let (+) t1 t2 = In(Op(Plus, t1, t2))
  let ( * ) t1 t2 = In(Op(Times, t1, t2))
  let (-) t1 t2 = In(Op(Minus, t1, t2))
  let ( .%[] ) t1 t2 = In(Select(t1, t2))
  let ( .%[]<- ) t1 t2 t3 = In(Store(t1, t2, t3))
end

type bop = Eq | Leq

type 'a propF = 
  | True 
  | And of 'a * 'a 
  | False
  | Or of 'a * 'a 
  | Not of 'a 
  | BOp of bop * term * term 
  | ArrayProp of var list * 'a * 'a  

let map_propF f (x : 'a propF) = 
  match x with
  | True -> True
  | And (a1, a2) -> And(f a1, f a2)
  | False -> False
  | Or (a1, a2) -> Or(f a1, f a2)
  | Not a1 -> Not (f a1)
  | BOp (op, t1, t2) -> BOp (op, t1, t2)
  | ArrayProp (xs, ig, vc) -> ArrayProp(xs, f ig, f vc)

type prop = In of prop propF 

module PropBuild = struct
  let true' = In True
  let false' = In False
  let ( && ) t1 t2 = In(And(t1, t2))
  let ( || ) t1 t2 = In(Or(t1, t2))
  let not t = In(Not t)
  let implies p1 p2 = not p1 || p2 

  let intersect ps = 
    match ps with
    | [] -> true'
    | p :: ps -> List.fold_left (&&) p ps 

  let bop op t1 t2 = In(BOp(op, t1, t2))

  let ( = ) t1 t2 = In(BOp(Eq, t1, t2))
  let ( <= ) t1 t2 = In(BOp(Leq, t1, t2))

  let ( >= ) t1 t2 = t2 <= t1 

  (* These are engineered to ensure that t1 is safe to be a uvar. If 
     we ever need i = j with both sides as uvars, then we need to 
     rewrite the formula. I'll ignore all that in this example. *)

  let ( < ) t1 t2 = let open TermBuild in 
                    t1 <= t2 - lit 1 

  let ( > ) t1 t2 = let open TermBuild in 
                    t2 - lit 1 <= t1 

  let (!=) t1 t2 = t1 < t2 || t1 > t2 
end

let rec gensym = 
  let r = ref 0 in
  fun s -> Printf.sprintf "%s_%d" s (incr r; !r) 

let rec rename_term xys ((In t) : term) : term = 
  In(match t with
     | Var x -> (match List.assoc_opt x xys with
                 | Some z -> Var z
                 | None -> Var x)
     | _ -> map_termF (rename_term xys) t)

let rec rename_prop xys ((In p) : prop) : prop = 
  In(match p with
     | BOp(op, t1, t2) -> BOp(op, rename_term xys t1, rename_term xys t2)
     | ArrayProp(ys, p1, p2) -> 
        let yzs = List.map (fun y -> (y, gensym y)) ys in
        let xys' = yzs @ xys in 
        ArrayProp(List.map snd yzs, 
                  rename_prop xys' p1, 
                  rename_prop xys' p2)
     | _ -> map_propF (rename_prop xys) p)

let rec subst_term env ((In t) : term) : term = 
  match t with
  | Var x -> (match List.assoc_opt x env with
              | Some t' -> t'
              | None -> In(Var x))
  | _ -> In(map_termF (subst_term env) t)

let subst_prop env p = 
  let rec loop env (In p) = 
    match p with
    | BOp (op, t1, t2) -> In(BOp(op, subst_term env t1, subst_term env t2))
    | _ -> In(map_propF (loop env) p)
  in
  loop env (rename_prop [] p)

let array_gensym ((In t) : term)  = 
  match t with
  | Var x -> gensym x 
  | _ -> gensym "a"

let rec replace_stores_in_term ((In t) : term) = 
  let open TermBuild in 
  let open PropBuild in 
  match t with
  | Var x -> var x, []
  | Lit n -> lit n, []
  | Op (op, t1, t2) -> let t1', r1 = replace_stores_in_term t1 in
                       let t2', r2 = replace_stores_in_term t2 in 
                       (In(Op(op, t1, t2))), (r1 @ r2)
  | Select (t1, t2) -> let t1', r1 = replace_stores_in_term t1 in
                       let t2', r2 = replace_stores_in_term t2 in
                       t1'.%[ t2'], (r1 @ r2)
  | Store (t1, t2, t3) -> 
     let t1', r1 = replace_stores_in_term t1 in
     let t2', r2 = replace_stores_in_term t2 in
     let t3', r3 = replace_stores_in_term t3 in
     let a = array_gensym t1 in 
     let p1 = (var a).%[ t2 ] = t3 in 
     let j = gensym "j" in 
     let p2 = In(ArrayProp([j], (var j <= t2 - lit 1) 
                             ||
                             (t2 + lit 1 <=  var j),
                             (var a).%[ var j ] = t1'.%[ var j]))
     in
     (var a, [p1 && p2] @ r1 @ r2 @ r3)

let rec replace_stores ((In p) : prop) = 
  let open PropBuild in 
  match p with
  | BOp (op, t1, t2) -> let t1', ps1 = replace_stores_in_term t1 in 
                        let t2', ps2= replace_stores_in_term t2 in 
                        let p = bop op t1' t2' in 
                        List.fold_left (&&) p (ps1 @ ps2)
  | _ -> In(map_propF replace_stores p)

let rec nnf ((In p) : prop) = 
  match p with
  | Not p1 -> nnf' p1 
  | _ -> In(map_propF nnf p)
and nnf' ((In p) : prop) : prop = 
  let open PropBuild in 
  match p with
  | True -> false'
  | And (p1, p2) -> nnf' p1 || nnf' p2
  | False -> true'
  | Or (p1, p2) -> nnf' p1 && nnf' p2
  | Not p1 -> nnf p1
  | BOp(Eq, t1, t2) -> t1 != t2 
  | BOp(Leq, t1, t2) -> t1 > t2 
  | ArrayProp (xs, p1, p2) -> 
     let xys = List.map (fun z -> (z, gensym z)) xs in
     let p1 = rename_prop xys p1 in
     let p2 = rename_prop xys p2 in 
     nnf p1 && nnf' p2

(* Compute all the non-uvar array indices *)

let read_set p = 
  let rec term uvars ((In t) : term) acc = 
    match t with
    | Var _ -> acc
    | Lit _ -> acc
    | Op (_, t1, t2) -> term uvars t2 (term uvars t1 acc)
    | Select (t1, In(Var x)) when List.mem x uvars -> term uvars t1 acc 
    | Select (t1, t2) -> term uvars t2 (term uvars t1 (t2 :: acc))
    | Store (_, _, _) -> assert false 
  in
  let rec prop uvars ((In p) : prop) acc = 
    match p with
    | True -> acc
    | And (p1, p2) -> prop uvars p2 (prop uvars p1 acc)
    | False -> acc
    | Or (p1, p2) -> prop uvars p2 (prop uvars p1 acc)
    | Not _ -> assert false 
    | BOp (_, t1, t2) -> term uvars t2 (term uvars t1 acc)
    | ArrayProp (xs, t1, t2) -> let uvars' = xs @ uvars in 
                                prop uvars' t2 (prop uvars' t1 acc)
  in
  prop [] p []
  
(* Compute all the guard expressions *)

let guard_set p = 
  let rec iguards uvars ((In p) : prop) acc = 
    match p with
    | True
    | False -> acc
    | And (p1, p2) 
    | Or (p1, p2) -> iguards uvars p2 (iguards uvars p1 acc)
    | Not _ 
    | ArrayProp (_, _, _) -> assert false (* iguard *)
    | BOp (op, t1, t2) -> 
       let is_uvar ((In t) : term) = 
         match t with
         | Var y when List.mem y uvars -> true
         | _ -> false
       in
       let new_guards = List.filter is_uvar [t1; t2] in 
       new_guards @ acc 
  in
  let rec prop uvars ((In p) : prop) acc = 
    match p with
    | True
    | False -> acc
    | And (p1, p2)
    | Or (p1, p2) -> prop uvars p2 (prop uvars p1 acc)
    | Not _ -> assert false (* NNF *)
    | BOp (_, _, _) -> acc
    | ArrayProp (xs, ig, bod) -> 
       let uvars' = xs @ uvars in 
       prop uvars' bod (iguards uvars' ig acc)
  in
  prop [] p []

let indices p = 
  let rs = read_set p in 
  let gs = guard_set p in 
  if rs = [] && gs = [] then 
    [TermBuild.lit 0]
  else
    rs @ gs 
                     
let rec replace_quantifers i_phi ((In p) : prop) = 
  let open PropBuild in 
  match p with
  | ArrayProp (xs, ig, body) -> (* Assume xs fresh wrt i_phi *)
     let q = implies ig (replace_quantifers i_phi body) in
     let rec envs xs = 
       match xs with 
       | [] -> [[]]
       | y :: ys -> 
          let envss = 
            List.map (fun env -> 
                List.map (fun t -> (y, t) :: env) i_phi)
              (envs ys)
          in
          List.concat envss
     in
     intersect (List.map (fun env -> subst_prop env q) (envs xs))
  | _ -> In(map_propF (replace_quantifers i_phi) p)

let normalise p0 = 
  let p1 = rename_prop [] p0 in 
  let p2 = replace_stores p1 in 
  let p3 = replace_quantifers (indices p2) p2 in 
  p3