Fixes primus self seeding generators (#964)

* memoizes values read from Primus memory layers

Memoizes values produced by static or dynamic layers, so that two
consequent reads will yield the same value.

* fixes Primus generators

Due to a design error, Primus self-seeding generators never worked as
expected, as they were recreating themselves on each access. The
problem is more a design error, rather than implementation, therefore
the fix is not very beatiful. It is really hard to implement correctly
the `create init` function which has `(int -> t) -> t` type. The proposed
implementation allows only such init that forall `x, y : int`, `init
x` and `init y` generate generators of the same type.

Probably, there is a better solution (by making the Generator type a
Primus component and turning them into higher-order functions), but
since we dropped generators in Primus 2.0 (in favor of a much easier to
use mechanism), I don't think that it is worthwhile to spend more time
on the current implementation.

* updates docs and error messages

lib/bap_primus/bap_primus.mli

@@ -1355,7 +1355,15 @@ module Std : sig
         module Seeded : sig
 
           (** [create init] creates a self-seeded generator from a
-              regular generator.  *)
+              regular generator.
+
+              Caveats:
+
+              The [init] function can use only one of the two
+              generator constructors to create a generator:
+              - [Random.lcg]
+              - [Random.byte]
+          *)
           val create : (int -> t) -> t
 
 

lib/bap_primus/bap_primus_generator.ml

@@ -17,37 +17,45 @@ let generators : Univ_map.t state = Bap_primus_machine.State.declare
     ~uuid:"7e81d5ae-46a2-42ff-918f-96c0c2dc95e3"
     (fun _ -> Univ_map.empty)
 
+module States = Univ_map
 
-module States = Univ_map.With_default
+module type Iter =
+  Iterator.Infinite.S with type dom = int
+
+type 'a iter = (module Iter with type t = 'a)
 
 type 'a gen = {
-  state : 'a States.Key.t;
-  next : 'a -> 'a;
-  value : 'a -> int;
-  min : int;
-  max : int;
+  iter : 'a iter;
+  self : 'a;
+}
+
+type 'a key = 'a gen States.Key.t
+
+type 'a ready = {
+  key : 'a key;
+  gen : 'a gen;
+}
+
+type 'a wait = {
+  key : 'a key;
+  init : int -> 'a gen
 }
 
 type t =
-  | Gen : 'a gen -> t
-  | Wait : (int -> t) -> t
   | Static : int -> t
+  | Ready : 'a ready -> t
+  | Wait : 'a wait -> t
 
 let rec sexp_of_t = function
   | Static x -> Sexp.(List [Atom "static"; sexp_of_int x])
-  | Gen {min;max} -> Sexp.List [
-      Sexp.Atom "interval";
-      sexp_of_int min;
-      sexp_of_int max;
-    ]
-  | Wait create -> Sexp.List [Sexp.Atom "project-0"; sexp_of_t (create 0)]
-
-let create (type rng)
-    (module Rng : Iterator.Infinite.S
-      with type t = rng and type dom = int) init =
+  | _ -> Sexp.Atom "<generator>"
+
+let make key init iter = Ready {key; gen={iter; self=init}}
+
+let create iter init =
   let state = States.Key.create
-      ~default:init ~name:"rng-state" sexp_of_opaque in
-  Gen {state; next=Rng.next; value=Rng.value; min=Rng.min; max=Rng.max}
+      ~name:"rng-state" sexp_of_opaque in
+  make state init iter
 
 let unfold (type gen)
     ?(min=Int.min_value)
@@ -70,32 +78,82 @@ let static value = Static value
 module Random = struct
   open Bap_primus_random
 
+  let lcg_key : (LCG.t * int) gen States.Key.t =
+    States.Key.create ~name:"linear-congruent-generator"
+      sexp_of_opaque
+
   let lcg ?(min=LCG.min) ?(max=LCG.max) seed =
-    let lcg = LCG.create seed in
-    unfold ~min ~max ~f:(fun (lcg,value) ->
-        let x = LCG.value lcg in
-        (LCG.next lcg, min + x mod (max-min+1))) lcg
+    let next (gen,_) =
+      let gen = LCG.next gen in
+      let x = min + LCG.value gen mod (max-min+1) in
+      gen,x in
+    let value = snd in
+    let init = next (LCG.create seed,0) in
+    make lcg_key init (module struct
+      type t = LCG.t * int
+      type dom = int
+      let min = min
+      let max = max
+      let next = next
+      let value = value
+    end)
 
   let byte seed = lcg ~min:0 ~max:255 seed
 
+  let cast_gen : type a b. a key -> b ready -> a gen option =
+    fun k1 {key=k2; gen} -> match States.Key.same_witness k1 k2 with
+      | None -> None
+      | Some Type_equal.T -> Some gen
+
+
   module Seeded = struct
-    let create make = Wait make
-    let lcg ?min ?max () = Wait (fun seed -> lcg ?min ?max seed)
-    let byte = Wait byte
+    let unpack_make key make =
+      let init seed  = match make seed with
+        | Wait _
+        | Static _ -> failwith "Generator.Seeded: invalid initializer"
+        | Ready g -> match cast_gen key g with
+          | Some g -> g
+          | None -> invalid_arg "Seeded.create changed its type" in
+      Wait {key; init}
+
+    let create make = match make 0 with
+      | Ready {key} -> unpack_make key make
+      | _ -> invalid_arg "Seeded.create must always create \
+                          an iterator of the same type"
+
+    let lcg ?min ?max () = unpack_make lcg_key (fun seed ->
+        lcg ?min ?max seed)
+
+    let byte = lcg ~min:0 ~max:255 ()
   end
 end
 
+
 module Make(Machine : Machine) = struct
   open Machine.Syntax
+
+  let call (type a) (state : a key) ({iter; self} : a gen) =
+    let module Iter : Iter with type t = a = (val iter) in
+    Machine.Local.get generators >>= fun states ->
+    let self = match States.find states state with
+      | None -> self
+      | Some {self} -> self in
+    let iter = {
+      iter;
+      self = Iter.next self;
+    } in
+    let states = States.set states state iter in
+    Machine.Local.put generators states >>| fun () ->
+    Iter.value iter.self
+
   let rec next = function
-    | Gen {state; next; value} ->
-      Machine.Local.get generators >>= fun states ->
-      let gen = States.find states state in
-      let states = States.set states state (next gen) in
-      Machine.Local.put generators states >>| fun () ->
-      value gen
-    | Wait create ->
-      Machine.current () >>= fun id ->
-      next (create (Machine.Id.hash id))
     | Static n -> Machine.return n
+    | Ready {key; gen} -> call key gen
+    | Wait {key; init} ->
+      Machine.Local.get generators >>= fun states ->
+      match States.find states key with
+      | None ->
+        Machine.current () >>= fun id ->
+        call key (init (Machine.Id.hash id))
+      | Some iter -> call key iter
 end

lib/bap_primus/bap_primus_memory.ml

@@ -103,16 +103,25 @@ module Make(Machine : Machine) = struct
 
   let pagefault addr = Machine.raise (Pagefault addr)
 
+  let remembered {values; layers} addr value =
+    Machine.Local.put state {
+      layers;
+      values = Map.set values ~key:addr ~data:value;
+    } >>| fun () ->
+    value
+
   let read addr {values;layers} = match find_layer addr layers with
     | None -> pagefault addr
     | Some layer -> match Map.find values addr with
       | Some v -> Machine.return v
-      | None -> match layer.mem with
-        | Dynamic {value} ->
-          Generate.next value >>= Value.of_int ~width:8
-        | Static mem -> match Memory.get ~addr mem with
-          | Ok v -> Value.of_word v
-          | Error _ -> failwith "Primus.Memory.read"
+      | None ->
+        let read_value = match layer.mem with
+          | Dynamic {value} ->
+            Generate.next value >>= Value.of_int ~width:8
+          | Static mem -> match Memory.get ~addr mem with
+            | Ok v -> Value.of_word v
+            | Error _ -> failwith "Primus.Memory.read" in
+        read_value >>= remembered {values; layers} addr
 
 
   let write addr value {values;layers} =