From 65fa3e84780448e6623400a27695b18f3392a237 Mon Sep 17 00:00:00 2001
From: zapashcanon <leo@ndrs.fr>
Date: Wed, 4 Sep 2024 18:26:49 +0200
Subject: [PATCH] small optimisation of the lazy memory model

---
 src/symbolic/symbolic_memory.ml              | 26 ++++++++++++++------
 src/symbolic/symbolic_memory_concretizing.ml | 26 +++++++++++++-------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/symbolic/symbolic_memory.ml b/src/symbolic/symbolic_memory.ml
index 69519e55e..dd9861c03 100644
--- a/src/symbolic/symbolic_memory.ml
+++ b/src/symbolic/symbolic_memory.ml
@@ -10,14 +10,14 @@ open Expr
 let page_size = Symbolic_value.const_i32 65_536l
 
 type t =
-  { data : (Int32.t, Value.int32) Hashtbl.t
+  { data : (Int32.t, Value.int32) Hashtbl.t option ref
   ; parent : t option
   ; mutable size : Value.int32
   ; chunks : (Int32.t, Value.int32) Hashtbl.t
   }
 
 let create size =
-  { data = Hashtbl.create 128
+  { data = ref None
   ; parent = None
   ; size = Value.const_i32 size
   ; chunks = Hashtbl.create 16
@@ -41,6 +41,14 @@ let fill _ = assert false
 
 let blit _ = assert false
 
+let replace m k v =
+  match !(m.data) with
+  | None ->
+    let tbl = Hashtbl.create 16 in
+    Hashtbl.add tbl k v;
+    m.data := Some tbl
+  | Some tbl -> Hashtbl.replace tbl k v
+
 let blit_string m str ~src ~dst ~len =
   (* This function is only used in memory init so everything will be concrete *)
   let str_len = String.length str in
@@ -54,20 +62,24 @@ let blit_string m str ~src ~dst ~len =
     for i = 0 to len - 1 do
       let byte = Char.code @@ String.get str (src + i) in
       let dst = Int32.of_int (dst + i) in
-      Hashtbl.replace m.data dst (make (Val (Num (I8 byte))))
+      replace m dst (make (Val (Num (I8 byte))))
     done;
     Value.Bool.const false
   end
 
 let clone m =
-  { data = Hashtbl.create 16
-  ; parent = Some m
+  let parent = if Option.is_none !(m.data) then m.parent else Some m in
+  { data = ref None
+  ; parent
   ; size = m.size
   ; chunks = Hashtbl.copy m.chunks (* TODO: we can make this lazy as well *)
   }
 
 let rec load_byte { parent; data; _ } a =
-  try Hashtbl.find data a
+  try
+    match !data with
+    | None -> raise Not_found
+    | Some data -> Hashtbl.find data a
   with Not_found -> (
     match parent with
     | None -> make (Val (Num (I8 0)))
@@ -161,7 +173,7 @@ let storen m ~addr v n =
   for i = 0 to n - 1 do
     let addr' = Int32.add a0 (Int32.of_int i) in
     let v' = extract v i in
-    Hashtbl.replace m.data addr' v'
+    replace m addr' v'
   done
 
 let store_8 m ~addr v = storen m ~addr v 1
diff --git a/src/symbolic/symbolic_memory_concretizing.ml b/src/symbolic/symbolic_memory_concretizing.ml
index 16fcd6724..4ab43aa02 100644
--- a/src/symbolic/symbolic_memory_concretizing.ml
+++ b/src/symbolic/symbolic_memory_concretizing.ml
@@ -4,20 +4,17 @@ module Backend = struct
   type address = Int32.t
 
   type t =
-    { data : (address, Symbolic_value.int32) Hashtbl.t
+    { data : (address, Symbolic_value.int32) Hashtbl.t option ref
     ; parent : t option
     ; chunks : (address, Symbolic_value.int32) Hashtbl.t
     }
 
-  let make () =
-    { data = Hashtbl.create 16; parent = None; chunks = Hashtbl.create 16 }
+  let make () = { data = ref None; parent = None; chunks = Hashtbl.create 16 }
 
   let clone m =
+    let parent = if Option.is_none !(m.data) then m.parent else Some m in
     (* TODO: Make chunk copying lazy *)
-    { data = Hashtbl.create 16
-    ; parent = Some m
-    ; chunks = Hashtbl.copy m.chunks
-    }
+    { data = ref None; parent; chunks = Hashtbl.copy m.chunks }
 
   let address a =
     let open Symbolic_choice_without_memory in
@@ -30,7 +27,10 @@ module Backend = struct
   let address_i32 a = a
 
   let rec load_byte { parent; data; _ } a =
-    try Hashtbl.find data a
+    try
+      match !data with
+      | None -> raise Not_found
+      | Some data -> Hashtbl.find data a
     with Not_found -> (
       match parent with
       | None -> Expr.value (Num (I8 0))
@@ -96,11 +96,19 @@ module Backend = struct
       sym
     | _ -> Expr.make (Extract (v, pos + 1, pos))
 
+  let replace m k v =
+    match !(m.data) with
+    | None ->
+      let tbl = Hashtbl.create 16 in
+      Hashtbl.add tbl k v;
+      m.data := Some tbl
+    | Some tbl -> Hashtbl.replace tbl k v
+
   let storen m a v n =
     for i = 0 to n - 1 do
       let a' = Int32.add a (Int32.of_int i) in
       let v' = extract v i in
-      Hashtbl.replace m.data a' v'
+      replace m a' v'
     done
 
   let validate_address m a range =