make memoization routines in Trans and Register simpler

src/core/trans.ml

@@ -43,57 +43,42 @@ let compose_l f g x = list_apply g (f x)
 
 let apply f x = f x
 
-let ymemo f tag h =
-  let rec aux x =
+let memo tag f =
+  let h = Hashtbl.create 63 in fun x ->
     let t = tag x in
-    try
-      Hashtbl.find h t
+    try Hashtbl.find h t
     with Not_found ->
-      let r = f aux x in
+      let r = f x in
       Hashtbl.add h t r;
-      r in
-  aux
-
-let memo f tag h = ymemo (fun _ -> f) tag h
+      r
 
 let term_tag t = t.t_tag
 let fmla_tag f = f.f_tag
 let decl_tag d = d.d_tag
 
-let task_tag = function
-  | Some t -> t.task_tag
-  | None   -> -1
-
-let store f = memo f task_tag (Hashtbl.create 63)
+let store f = memo task_tag f
 
-let accum memo_t rewind v =
+let fold fn v =
+  let h = Hashtbl.create 63 in
+  let rewind acc task =
+    let acc = fn task acc in
+    Hashtbl.add h task.task_tag acc;
+    acc
+  in
+  let curr task =
+    try Some (Hashtbl.find h task.task_tag)
+    with Not_found -> None
+  in
   let rec accum todo = function
-    | Some task ->
-        let curr =
-          try Some (Hashtbl.find memo_t task.task_tag)
-          with Not_found -> None
-        in
-        begin match curr with
-          | Some curr -> rewind curr todo
-          | None      -> accum (task::todo) task.task_prev
-        end
-    | None -> rewind v todo
+    | None -> List.fold_left rewind v todo
+    | Some task -> begin match curr task with
+        | Some v -> List.fold_left rewind v todo
+        | None   -> accum (task::todo) task.task_prev
+      end
   in
-  accum
-
-let save memo_t task v = Hashtbl.add memo_t task.task_tag v; v
-
-let fold fn v =
-  let memo_t = Hashtbl.create 63 in
-  let rewind x task = save memo_t task (fn task x) in
-  let rewind = List.fold_left rewind in
-  accum memo_t rewind v []
+  accum []
 
-let fold_l fn v =
-  let memo_t = Hashtbl.create 63 in
-  let rewind x task = save memo_t task (list_apply (fn task) x) in
-  let rewind = List.fold_left rewind in
-  accum memo_t rewind [v] []
+let fold_l fn v = fold (fun task -> list_apply (fn task)) [v]
 
 let fold_map   fn v t = conv_res snd                (fold   fn (v, t))
 let fold_map_l fn v t = conv_res (List.rev_map snd) (fold_l fn (v, t))
@@ -102,8 +87,7 @@ let map   fn = fold   (fun t1 t2 -> fn t1 t2)
 let map_l fn = fold_l (fun t1 t2 -> fn t1 t2)
 
 let decl fn =
-  let memo_t = Hashtbl.create 63 in
-  let fn d = memo fn decl_tag memo_t d in
+  let fn = memo decl_tag fn in
   let fn task acc = match task.task_decl with
     | Decl d -> List.fold_left add_decl acc (fn d)
     | td -> add_tdecl acc td
@@ -111,8 +95,7 @@ let decl fn =
   map fn
 
 let decl_l fn =
-  let memo_t = Hashtbl.create 63 in
-  let fn d = memo fn decl_tag memo_t d in
+  let fn = memo decl_tag fn in
   let fn task acc = match task.task_decl with
     | Decl d -> List.rev_map (List.fold_left add_decl acc) (fn d)
     | td -> [add_tdecl acc td]

src/driver/register.ml

@@ -41,35 +41,28 @@ let create gen = {
 
 exception ArgumentNeeded
 
-let memo tag f =
-  let h = Hashtbl.create 7 in
-  function
-  | None -> raise ArgumentNeeded
-  | Some x ->
-      let t = tag x in
-      try Hashtbl.find h t
-      with Not_found ->
-        let r = f x in
-        Hashtbl.add h t r;
-        r
-
-let memo_env e   = memo Env.env_tag e
-let memo_query q = memo query_tag q
-
-let memo2 extract f =
-  let h = Hashtbl.create 7 in
-  fun x ->
-    let arg,tag = extract x in
-    try Hashtbl.find h tag
+let memo tag ext f =
+  let h = Hashtbl.create 7 in fun x ->
+    let t = tag x in
+    try Hashtbl.find h t
     with Not_found ->
-      let r = f arg x in
-      Hashtbl.add h tag r; r
+      let r = f (ext x) in
+      Hashtbl.add h t r;
+      r
 
-let memo_use x   = memo2 (fun t -> task_used t,  task_tag (last_use   t)) x
-let memo_clone x = memo2 (fun t -> task_clone t, task_tag (last_clone t)) x
-let memo_goal x  = memo2 (fun t -> t, task_tag t) x
+let memo_env f =
+  let f = memo Env.env_tag (fun e -> e) f in function
+    | None -> raise ArgumentNeeded
+    | Some env -> f env
 
-let query drv task = Util.option_map (fun d -> driver_query d task) drv
+let memo_query f =
+  let f = memo query_tag (fun q -> q) f in function
+    | None -> raise ArgumentNeeded
+    | Some drv -> fun task -> f (driver_query drv task) task
+
+let memo_use   f = memo (fun t -> task_tag (last_use   t)) task_used  f
+let memo_clone f = memo (fun t -> task_tag (last_clone t)) task_clone f
+let memo_goal  f = memo task_tag (fun t -> t) f
 
 let gen_gen f () =
   let f0 = Trans.apply (f ()) in
@@ -78,31 +71,31 @@ let gen_gen f () =
 let gen_env f () =
   let f0 env = Trans.apply (f env) in
   let f1 = memo_env f0 in
-  fun env _ -> f1 env
+  fun env _ task -> f1 env task
 
 let gen_query f () =
   let f0 query = Trans.apply (f query) in
   let f1 = memo_query f0 in
-  fun _ drv task -> f1 (query drv task) task
+  fun _ drv task -> f1 drv task
 
 let gen_arg memo_arg f () =
   let f0 env arg = Trans.apply (f env arg) in
   let f1 env = memo_arg (f0 env) in
   let f2 = memo_env f1 in
-  fun env _ -> f2 env
+  fun env _ task -> f2 env task task
 
 let gen_full f () =
   let f0 query goal = Trans.apply (f query goal) in
   let f1 query = memo_goal (f0 query) in
   let f2 = memo_query f1 in
-  fun _ drv task -> f2 (query drv task) task
+  fun _ drv task -> f2 drv task task
 
 let gen_both f () =
   let f0 env use clone = Trans.apply (f env use clone) in
   let f1 env use = memo_clone (f0 env use) in
   let f2 env = memo_use (f1 env) in
   let f3 = memo_env f2 in
-  fun env _ -> f3 env
+  fun env _ task -> f3 env task task task
 
 let store       f = create (gen_gen f)
 let store_env   f = create (gen_env f)