Use Patricia trees

dune-project

@@ -37,6 +37,7 @@
   eio_main
   simple_smt
   iter
+  patricia-tree
   (ppx_expect
    (>= 0.17.0))
   ppx_deriving
@@ -144,6 +145,7 @@
   soteria
   dune-site
   zarith
+  patricia-tree
   ppx_deriving
   cmdliner
   fmt

soteria-c/lib/csymex.ml

@@ -62,8 +62,11 @@ module With_origin =
     end)
 
 module Freeable = Soteria.Sym_states.Freeable.Make (SYMEX)
-module Pmap_direct_access = Soteria.Sym_states.Pmap.Direct_access (SYMEX)
-module Pmap = Soteria.Sym_states.Pmap.Make (SYMEX)
+
+module Pmap_direct_access =
+  Soteria.Sym_states.Pmap.Direct_access_patricia_tree (SYMEX)
+
+module Pmap = Soteria.Sym_states.Pmap.Make_patricia_tree (SYMEX)
 module Tree_block = Soteria.Sym_states.Tree_block.Make (SYMEX)
 module Concrete_map = Soteria.Sym_states.Pmap.Concrete (SYMEX)
 module Bi = Soteria.Sym_states.Bi_abd.Make (SYMEX)

soteria-c/lib/fun_ctx.ml

@@ -1,19 +1,9 @@
 module Bidirectional_map = struct
-  module SMap = Symbol_std.Map
-  module IMap = Map.Make (Z)
+  include Bimap.Make (Symbol_std) (Z)
 
-  type t = { s_to_i : Z.t SMap.t; i_to_s : Symbol_std.t IMap.t }
-
-  let empty = { s_to_i = SMap.empty; i_to_s = IMap.empty }
-  let has_sym s map = SMap.mem s map.s_to_i
-
-  let add s i map =
-    let s_to_i = SMap.add s i map.s_to_i in
-    let i_to_s = IMap.add i s map.i_to_s in
-    { s_to_i; i_to_s }
-
-  let get_loc_id s map = SMap.find_opt s map.s_to_i
-  let get_sym i map = IMap.find_opt i map.i_to_s
+  let has_sym = mem_l
+  let get_loc_id = find_l
+  let get_sym = find_r
 end
 
 (* FIXME: This is slightly off because we could have a location that is already assigned.

soteria-c/lib/globs.ml

@@ -42,24 +42,22 @@ let get sym st =
            It'd be nice to enforce this modularly, but right now we do it as a wrapper.
            Unless one has hundreds of globals, the cost is negligeable, but still. *)
   let open Csymex.Syntax in
-  let existed = Option.fold ~none:false ~some:(Sym_map.M.mem sym) st in
+  let existed = Option.fold ~none:false ~some:(Sym_map.syntactic_mem sym) st in
   let* res = (Sym_map.wrap GlobFn.load) sym st in
   let loc, st = Soteria.Symex.Compo_res.get_ok res in
   let+ () =
     if existed then (* We haven't created a new location *) return ()
     else
       (* We learn that the new location is distinct from all other global locations *)
-      let to_assume =
-        Sym_map.M.to_seq (Option.value ~default:Sym_map.M.empty st)
-        |> Seq.filter_map (fun (k, v) ->
-               let open Typed.Infix in
-               if Cerb_frontend.Symbol.equal_sym k sym then None
-               else
-                 let neq = Typed.not (v ==@ loc) in
-                 Some neq)
-        |> List.of_seq
-      in
-      assume to_assume
+      Sym_map.syntactic_bindings (Option.value ~default:Sym_map.empty st)
+      |> Seq.filter_map (fun (k, v) ->
+             let open Typed.Infix in
+             if Cerb_frontend.Symbol.equal_sym k sym then None
+             else
+               let neq = Typed.not (v ==@ loc) in
+               Some neq)
+      |> List.of_seq
+      |> assume
   in
   (loc, st)
 
@@ -69,14 +67,12 @@ let produce serialized st =
   let+ () =
     (* Bit heavy-handed but we just massively assume the well-formedness *)
     let to_assume =
-      Sym_map.M.to_seq (Option.value ~default:Sym_map.M.empty st')
-      |> Seq.self_cross_product
-      |> Seq.map (fun ((_, loca), (_, locb)) ->
-             let open Typed.Infix in
-             Typed.not (loca ==@ locb))
+      Sym_map.syntactic_bindings (Option.value ~default:Sym_map.empty st')
+      |> Seq.map snd
       |> List.of_seq
+      |> Typed.distinct
     in
-    assume to_assume
+    assume [ to_assume ]
   in
   st'
 

soteria-c/lib/state.ml

@@ -21,6 +21,7 @@ module SPmap = Pmap_direct_access (struct
   let pp = ppa
   let fresh () = Csymex.nondet Typed.t_loc
   let simplify = Csymex.simplify
+  let to_int = unique_tag
 end)
 
 type t = (Block.t SPmap.t option, Globs.t) State_intf.Template.t

soteria-c/scripts/experiments.py

@@ -249,7 +249,7 @@ def run_soteria_c(self):
                     "Hyperfine is required for benchmarking but was not found in PATH. Install it and try again."
                 )
                 exit(4)
-            cmd = f"hyperfine '{cmd}' --warmup 1 --runs 10"
+            cmd = f"hyperfine '{cmd}' --warmup 1 --runs 10 -i"
         self.run_command(cmd)
 
     def run(self):

soteria-rust.opam

@@ -19,6 +19,7 @@ depends: [
   "soteria"
   "dune-site"
   "zarith"
+  "patricia-tree"
   "ppx_deriving"
   "cmdliner"
   "fmt"

soteria-rust/lib/config.ml

@@ -2,7 +2,9 @@
    We ignore the warning here. *)
 [@@@warning "-unused-open"]
 
-type frontend = Charon | Obol [@@deriving subliner_enum]
+type frontend = Charon | Obol
+[@@deriving subliner_enum, show { with_path = false }]
+
 type provenance = Strict | Permissive [@@deriving subliner_enum]
 
 type t = {

soteria-rust/lib/dune

@@ -4,18 +4,19 @@
  (package soteria-rust)
  (name soteria_rust_lib)
  (libraries
-  soteria
+  charon
   cmdliner
+  dune-site
   fmt
   fmt.tty
-  iter
-  zarith
   grace
   grace.ansi_renderer
+  iter
   simple_smt
-  charon
-  dune-site
-  str)
+  soteria
+  str
+  patricia-tree
+  zarith)
  (flags :standard -open Soteria.Soteria_std -open Soteria.Bv_values)
  (preprocess
   (pps ppx_deriving.std ppx_subliner soteria.ppx))

soteria-rust/lib/frontend.ml

@@ -427,13 +427,18 @@ let parse_ullbc ~mode ~plugin ?input ~output ~pwd () =
           err;
     Cleaner.touched output);
   let crate =
-    try
+    match
       output |> Yojson.Basic.from_file |> Charon.UllbcOfJson.crate_of_json
     with
-    | Sys_error _ -> frontend_err "File doesn't exist"
-    | _ -> frontend_err "Failed to parse ULLBC"
+    | Ok crate -> crate
+    | Error _ ->
+        Fmt.kstr frontend_err
+          "Failed to parse ULLBC. Do you have the right version of %a \
+           installed?"
+          Config.pp_frontend !Config.current.frontend
+    | exception Sys_error _ -> frontend_err "File doesn't exist"
+    | exception _ -> frontend_err "Unexpected error"
   in
-  let crate = Result.get_or_raise frontend_err crate in
   if !Config.current.output_crate then (
     (* save pretty-printed crate to local file *)
     let crate_file = Printf.sprintf "%s.crate" output in

soteria-rust/lib/sptr.ml

@@ -46,12 +46,14 @@ module DecayMap : DecayMapS = struct
 
     type t = T.sloc Typed.t
 
+    let to_int = unique_tag
     let pp = ppa
     let simplify = Rustsymex.simplify
     let fresh () = failwith "Allocation is not valid for the decay map!"
   end
 
-  module SPmap = Soteria.Sym_states.Pmap.Make (Rustsymex) (StateKey)
+  module SPmap =
+    Soteria.Sym_states.Pmap.Make_patricia_tree (Rustsymex) (StateKey)
 
   type data = { address : T.sint Typed.t; exposed : bool }
   [@@deriving show { with_path = false }]
@@ -95,13 +97,13 @@ module DecayMap : DecayMapS = struct
        See https://doc.rust-lang.org/nightly/std/ptr/fn.with_exposed_provenance.html
        *)
     let usize_ty = Typed.t_usize () in
-    let bindings = SPmap.M.bindings (SPmap.of_opt st) in
+    let bindings = SPmap.syntactic_bindings (SPmap.of_opt st) in
     let binding =
       Typed.iter_vars loc_int
       |> Iter.filter (fun (_, ty) -> Typed.equal_ty usize_ty ty)
       |> Iter.filter_map (fun (var, ty) ->
              let v = Typed.mk_var var ty in
-             List.find_opt
+             Seq.find
                (fun (_, { address; exposed }) ->
                  exposed && Typed.equal v address)
                bindings)

soteria-rust/lib/state.ml

@@ -44,6 +44,7 @@ module StateKey = struct
   type t = T.sloc Typed.t
 
   let pp = ppa
+  let to_int = unique_tag
   let i = ref 0
 
   let fresh_rsym () =
@@ -59,7 +60,10 @@ module StateKey = struct
 end
 
 module Freeable = Soteria.Sym_states.Freeable.Make (DecayMapMonad)
-module SPmap = Soteria.Sym_states.Pmap.Direct_access (DecayMapMonad) (StateKey)
+
+module SPmap =
+  Soteria.Sym_states.Pmap.Direct_access_patricia_tree (DecayMapMonad) (StateKey)
+
 module Bi = Soteria.Sym_states.Bi_abd.Make (DecayMapMonad)
 module Tree_block = Rtree_block.Make (Sptr)
 
@@ -85,7 +89,7 @@ end)
 
 module FunBiMap = struct
   include
-    Bimap.Make
+    Bimap.MakePp
       (struct
         type t = T.sloc Typed.t
 

soteria-rust/lib/tree_borrow.ml

@@ -11,10 +11,10 @@ let fresh_tag () =
 let zero = 0
 let pp_tag fmt tag = Fmt.pf fmt "‖%d‖" tag
 
-module TagMap = Map.Make (struct
+module TagMap = PatriciaTree.MakeMap (struct
   type t = tag
 
-  let compare = compare
+  let to_int = Fun.id
 end)
 
 type access = Read | Write
@@ -41,7 +41,7 @@ let pp_state fmt = function
   | Disabled -> Fmt.string fmt "Dis "
   | UB -> Fmt.string fmt "UB  "
 
-let meet st1 st2 =
+let[@inline] meet st1 st2 =
   match (st1, st2) with
   | UB, _ | _, UB -> UB
   | Disabled, _ | _, Disabled -> Disabled
@@ -52,7 +52,7 @@ let meet st1 st2 =
   | Reserved _, ReservedIM | ReservedIM, Reserved _ ->
       failwith "Can't compare Reserved and ReservedIM"
 
-let meet' (p1, st1) (p2, st2) = (p1 || p2, meet st1 st2)
+let[@inline] meet' (p1, st1) (p2, st2) = (p1 || p2, meet st1 st2)
 
 let transition =
   let transition st e =
@@ -185,4 +185,4 @@ let access accessed e (root : t) st =
   in
   if !ub_happened then Result.error `AliasingError else Result.ok st'
 
-let merge = TagMap.merge @@ fun _ -> Option.merge meet'
+let merge = TagMap.idempotent_union @@ fun _ -> meet'

soteria.opam

@@ -19,6 +19,7 @@ depends: [
   "eio_main"
   "simple_smt"
   "iter"
+  "patricia-tree"
   "ppx_expect" {>= "0.17.0"}
   "ppx_deriving"
   "fmt"

soteria/lib/bv_values/analyses.ml

@@ -243,8 +243,8 @@ module Interval : S = struct
   let get n v st =
     match Var.Map.find_opt v st with Some r -> r | None -> Data.mk n
 
-  let st_equal = Var.Map.equal Data.equal
-  let merge_states = Var.Map.merge (fun _ -> Option.merge Data.union)
+  let st_equal = Var.Map.reflexive_equal Data.equal
+  let merge_states = Var.Map.idempotent_inter (fun _ -> Data.union)
 
   let pp ft st =
     Fmt.(iter_bindings Var.Map.iter (pair ~sep:(any " -> ") Var.pp Data.pp))
@@ -503,12 +503,17 @@ end
 
 module Equality : S = struct
   module UnionFind = UnionFind.Make (UnionFind.StoreMap)
-  module IMap = Map.Make (Int)
+
+  module VMap = PatriciaTree.MakeMap (struct
+    type t = Svalue.t
+
+    let to_int = Svalue.unique_tag
+  end)
 
   include Reversible.Make_mutable (struct
-    type t = Svalue.t UnionFind.store * Svalue.t UnionFind.rref IMap.t
+    type t = Svalue.t UnionFind.store * Svalue.t UnionFind.rref VMap.t
 
-    let default = (UnionFind.new_store (), IMap.empty)
+    let default = (UnionFind.new_store (), VMap.empty)
   end)
 
   (* override to account for UnionFind  *)
@@ -538,11 +543,11 @@ module Equality : S = struct
     | BvOfFloat _ | FloatOfBv _ -> 4
     | _ -> 1
 
-  let get_or_make (v : Svalue.t) ((uf, refs) as st) =
-    match IMap.find_opt v.tag refs with
+  let get_or_make v ((uf, refs) as st) =
+    match VMap.find_opt v refs with
     | None ->
         let ref = UnionFind.make uf v in
-        let refs = IMap.add v.tag ref refs in
+        let refs = VMap.add v ref refs in
         (ref, (uf, refs))
     | Some ref -> (ref, st)
 
@@ -552,18 +557,18 @@ module Equality : S = struct
          (fun v1 v2 -> if cost v1 > cost v2 then v2 else v1)
          v1 v2
 
-  let find_cheaper_opt (v : Svalue.t) (uf, refs) =
-    Option.bind (IMap.find_opt v.tag refs) @@ fun r ->
+  let find_cheaper_opt v (uf, refs) =
+    Option.bind (VMap.find_opt v refs) @@ fun r ->
     let v_repr = UnionFind.get uf r in
     if Svalue.equal v v_repr then None else Some v_repr
 
-  let known_eq (v1 : Svalue.t) (v2 : Svalue.t) (uf, refs) : bool =
-    match (IMap.find_opt v1.tag refs, IMap.find_opt v2.tag refs) with
+  let known_eq v1 v2 (uf, refs) : bool =
+    match (VMap.find_opt v1 refs, VMap.find_opt v2 refs) with
     | Some r1, Some r2 -> UnionFind.eq uf r1 r2
     | _ -> false
 
-  let eval_var (uf, refs) (var : Svalue.t) _ _ =
-    IMap.find_opt var.tag refs |> Option.fold ~none:var ~some:(UnionFind.get uf)
+  let eval_var (uf, refs) var _ _ =
+    VMap.find_opt var refs |> Option.fold ~none:var ~some:(UnionFind.get uf)
 
   let simplify (v : Svalue.t) st =
     let rec simplify ~fuel v =
@@ -595,6 +600,10 @@ module Equality : S = struct
         let v1, st = get_or_make v1 st in
         let v2, st = get_or_make v2 st in
         merge v1 v2 st;
+        (* Here we choose to pass down the equality (rather than simplifying to true),
+           so that the underlying solver can do trivial SAT checks using these equalities.
+           This also means this analysis is exempt from implementing [encode], since it
+           doesn't store anything. *)
         ((v, Var.Set.empty), st)
     | _ -> ((v, Var.Set.empty), st)
 

soteria/lib/bv_values/svalue.ml

@@ -207,6 +207,7 @@ and t = t_node hash_consed [@@deriving show { with_path = false }, eq, ord]
 
 let hash t = t.tag
 let kind t = t.node.kind
+let unique_tag t = t.tag
 
 let iter =
   let rec aux (f : t -> unit) (sv : t) : unit =

soteria/lib/bv_values/typed.mli

@@ -89,6 +89,7 @@ val ppa : Format.formatter -> 'a t -> unit
 val equal : ([< any ] as 'a) t -> 'a t -> bool
 val compare : ([< any ] as 'a) t -> 'a t -> int
 val hash : [< any ] t -> int
+val unique_tag : [< any ] t -> int
 
 (** Typed constructors *)