From 16d85543842c9eb91a4a0c21a9986926f4ff0b10 Mon Sep 17 00:00:00 2001
From: Denis Merigoux <denis.merigoux@gmail.com>
Date: Wed, 12 Jan 2022 16:25:46 +0100
Subject: [PATCH] Provide a function that removes all log calls

---
 compiler/dcalc/optimizations.ml  | 51 +++++++++++++++++++++++++++++++-
 compiler/dcalc/optimizations.mli |  2 ++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/compiler/dcalc/optimizations.ml b/compiler/dcalc/optimizations.ml
index 0149ed0b..9193476a 100644
--- a/compiler/dcalc/optimizations.ml
+++ b/compiler/dcalc/optimizations.ml
@@ -77,7 +77,7 @@ let rec partial_evaluation (ctx : partial_evaluation_ctx) (e : expr Pos.marked)
   | ELit l -> Bindlib.box (ELit l, pos)
   | EAbs ((binder, binder_pos), typs) ->
       let vars, body = Bindlib.unmbind binder in
-      let new_body = partial_evaluation ctx body in
+      let new_body = rec_helper body in
       let new_binder = Bindlib.bind_mvar vars new_body in
       Bindlib.box_apply (fun binder -> (EAbs ((binder, binder_pos), typs), pos)) new_binder
   | EApp (f, args) ->
@@ -144,3 +144,52 @@ let program_map (t : 'a -> expr Pos.marked -> expr Pos.marked Bindlib.box) (ctx
   }
 
 let optimize_program (p : program) : program = program_map partial_evaluation VarMap.empty p
+
+let rec remove_all_logs (e : expr Pos.marked) : expr Pos.marked Bindlib.box =
+  let pos = Pos.get_position e in
+  let rec_helper = remove_all_logs in
+  match Pos.unmark e with
+  | EVar (x, _) -> Bindlib.box_apply (fun x -> (x, pos)) (Bindlib.box_var x)
+  | ETuple (args, s_name) ->
+      Bindlib.box_apply
+        (fun args -> (ETuple (args, s_name), pos))
+        (List.map rec_helper args |> Bindlib.box_list)
+  | ETupleAccess (arg, i, s_name, typs) ->
+      Bindlib.box_apply (fun arg -> (ETupleAccess (arg, i, s_name, typs), pos)) (rec_helper arg)
+  | EInj (arg, i, e_name, typs) ->
+      Bindlib.box_apply (fun arg -> (EInj (arg, i, e_name, typs), pos)) (rec_helper arg)
+  | EMatch (arg, arms, e_name) ->
+      Bindlib.box_apply2
+        (fun arg arms -> (EMatch (arg, arms, e_name), pos))
+        (rec_helper arg)
+        (List.map rec_helper arms |> Bindlib.box_list)
+  | EArray args ->
+      Bindlib.box_apply
+        (fun args -> (EArray args, pos))
+        (List.map rec_helper args |> Bindlib.box_list)
+  | ELit l -> Bindlib.box (ELit l, pos)
+  | EAbs ((binder, binder_pos), typs) ->
+      let vars, body = Bindlib.unmbind binder in
+      let new_body = rec_helper body in
+      let new_binder = Bindlib.bind_mvar vars new_body in
+      Bindlib.box_apply (fun binder -> (EAbs ((binder, binder_pos), typs), pos)) new_binder
+  | EApp (f, args) ->
+      Bindlib.box_apply2
+        (fun f args ->
+          match (Pos.unmark f, args) with
+          | EOp (Unop (Log _)), [ arg ] -> arg
+          | _ -> (EApp (f, args), pos))
+        (rec_helper f)
+        (List.map rec_helper args |> Bindlib.box_list)
+  | EAssert e1 -> Bindlib.box_apply (fun e1 -> (EAssert e1, pos)) (rec_helper e1)
+  | EOp op -> Bindlib.box (EOp op, pos)
+  | EDefault (exceptions, just, cons) ->
+      Bindlib.box_apply3
+        (fun exceptions just cons -> (EDefault (exceptions, just, cons), pos))
+        (List.map rec_helper exceptions |> Bindlib.box_list)
+        (rec_helper just) (rec_helper cons)
+  | EIfThenElse (e1, e2, e3) ->
+      Bindlib.box_apply3
+        (fun e1 e2 e3 -> (EIfThenElse (e1, e2, e3), pos))
+        (rec_helper e1) (rec_helper e2) (rec_helper e3)
+  | ErrorOnEmpty e1 -> Bindlib.box_apply (fun e1 -> (ErrorOnEmpty e1, pos)) (rec_helper e1)
diff --git a/compiler/dcalc/optimizations.mli b/compiler/dcalc/optimizations.mli
index 50ba2b92..7ba886a2 100644
--- a/compiler/dcalc/optimizations.mli
+++ b/compiler/dcalc/optimizations.mli
@@ -18,3 +18,5 @@ open Ast
 val optimize_expr : expr Pos.marked -> expr Pos.marked Bindlib.box
 
 val optimize_program : program -> program
+
+val remove_all_logs : expr Pos.marked -> expr Pos.marked