diff --git a/lib/bap_primus/bap_primus.mli b/lib/bap_primus/bap_primus.mli index 8d541ecfa..575e87241 100644 --- a/lib/bap_primus/bap_primus.mli +++ b/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 diff --git a/lib/bap_primus/bap_primus_generator.ml b/lib/bap_primus/bap_primus_generator.ml index 92cacc662..af406c9b3 100644 --- a/lib/bap_primus/bap_primus_generator.ml +++ b/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 "" + +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 diff --git a/lib/bap_primus/bap_primus_memory.ml b/lib/bap_primus/bap_primus_memory.ml index 578182407..7c4e49369 100644 --- a/lib/bap_primus/bap_primus_memory.ml +++ b/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} =