From 1e1b0f352d722329333f7dffc51c65730ce3af23 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 30 Jul 2022 21:17:55 -0400
Subject: [PATCH] feat(simplify): add sidekick_simplify library

---
 src/simplify/dune                  |  6 +++
 src/simplify/sidekick_simplify.ml  | 80 ++++++++++++++++++++++++++++++
 src/simplify/sidekick_simplify.mli | 39 +++++++++++++++
 3 files changed, 125 insertions(+)
 create mode 100644 src/simplify/dune
 create mode 100644 src/simplify/sidekick_simplify.ml
 create mode 100644 src/simplify/sidekick_simplify.mli

diff --git a/src/simplify/dune b/src/simplify/dune
new file mode 100644
index 00000000..fa4a9b50
--- /dev/null
+++ b/src/simplify/dune
@@ -0,0 +1,6 @@
+(library
+ (name Sidekick_simplify)
+ (public_name sidekick.simplify)
+ (synopsis "Simplifier")
+ (libraries containers iter sidekick.core sidekick.util)
+ (flags :standard -w +32 -open Sidekick_util))
diff --git a/src/simplify/sidekick_simplify.ml b/src/simplify/sidekick_simplify.ml
new file mode 100644
index 00000000..2daa4114
--- /dev/null
+++ b/src/simplify/sidekick_simplify.ml
@@ -0,0 +1,80 @@
+open Sidekick_core
+
+open struct
+  module P = Proof_trace
+  module Rule_ = Proof_core
+end
+
+type t = {
+  tst: Term.store;
+  proof: Proof_trace.t;
+  mutable hooks: hook list;
+  (* store [t --> u by step_ids] in the cache.
+     We use a bag for the proof steps because it gives us structural
+     sharing of subproofs. *)
+  cache: (Term.t * Proof_step.id Bag.t) Term.Tbl.t;
+}
+
+and hook = t -> Term.t -> (Term.t * Proof_step.id Iter.t) option
+
+let create tst ~proof : t =
+  { tst; proof; hooks = []; cache = Term.Tbl.create 32 }
+
+let[@inline] tst self = self.tst
+let[@inline] proof self = self.proof
+let add_hook self f = self.hooks <- f :: self.hooks
+let clear self = Term.Tbl.clear self.cache
+
+let normalize (self : t) (t : Term.t) : (Term.t * Proof_step.id) option =
+  (* compute and cache normal form of [t] *)
+  let rec loop t : Term.t * _ Bag.t =
+    match Term.Tbl.find self.cache t with
+    | res -> res
+    | exception Not_found ->
+      let steps_u = ref Bag.empty in
+      let u = aux_rec ~steps:steps_u t self.hooks in
+      Term.Tbl.add self.cache t (u, !steps_u);
+      u, !steps_u
+  and loop_add ~steps t =
+    let u, pr_u = loop t in
+    steps := Bag.append !steps pr_u;
+    u
+  (* try each function in [hooks] successively, and rewrite subterms *)
+  and aux_rec ~steps t hooks : Term.t =
+    match hooks with
+    | [] ->
+      let u =
+        Term.map_shallow self.tst ~f:(fun _inb u -> loop_add ~steps u) t
+      in
+      if Term.equal t u then
+        t
+      else
+        loop_add ~steps u
+    | h :: hooks_tl ->
+      (match h self t with
+      | None -> aux_rec ~steps t hooks_tl
+      | Some (u, _) when Term.equal t u -> aux_rec ~steps t hooks_tl
+      | Some (u, pr_u) ->
+        let bag_u = Bag.of_iter pr_u in
+        steps := Bag.append !steps bag_u;
+        let v, pr_v = loop u in
+        (* fixpoint *)
+        steps := Bag.append !steps pr_v;
+        v)
+  in
+  let u, pr_u = loop t in
+  if Term.equal t u then
+    None
+  else (
+    (* proof: [sub_proofs |- t=u] by CC + subproof *)
+    let step =
+      P.add_step self.proof
+      @@ Rule_.lemma_preprocess t u ~using:(Bag.to_iter pr_u)
+    in
+    Some (u, step)
+  )
+
+let normalize_t self t =
+  match normalize self t with
+  | Some (u, pr_u) -> u, Some pr_u
+  | None -> t, None
diff --git a/src/simplify/sidekick_simplify.mli b/src/simplify/sidekick_simplify.mli
new file mode 100644
index 00000000..1c3abf1a
--- /dev/null
+++ b/src/simplify/sidekick_simplify.mli
@@ -0,0 +1,39 @@
+(** Term simplifier *)
+
+open Sidekick_core
+
+type t
+
+val tst : t -> Term.store
+
+val create : Term.store -> proof:Proof_trace.t -> t
+(** Create a simplifier *)
+
+val clear : t -> unit
+(** Reset internal cache, etc. *)
+
+val proof : t -> Proof_trace.t
+(** Access proof *)
+
+type hook = t -> Term.t -> (Term.t * Proof_step.id Iter.t) option
+(** Given a Term.t, try to simplify it. Return [None] if it didn't change.
+
+        A simple example could be a hook that takes a Term.t [t],
+        and if [t] is [app "+" (const x) (const y)] where [x] and [y] are number,
+        returns [Some (const (x+y))], and [None] otherwise.
+
+        The simplifier will take care of simplifying the resulting Term.t further,
+        caching (so that work is not duplicated in subterms), etc.
+    *)
+
+val add_hook : t -> hook -> unit
+
+val normalize : t -> Term.t -> (Term.t * Proof_step.id) option
+(** Normalize a Term.t using all the hooks. This performs
+    a fixpoint, i.e. it only stops when no hook applies anywhere inside
+    the Term.t. *)
+
+val normalize_t : t -> Term.t -> Term.t * Proof_step.id option
+(** Normalize a Term.t using all the hooks, along with a proof that the
+    simplification is correct.
+    returns [t, ø] if no simplification occurred. *)