From 1bb338526d27bc7144836a0025fbd0f3cc8cab5a Mon Sep 17 00:00:00 2001
From: Denis Merigoux <denis.merigoux@gmail.com>
Date: Fri, 21 Apr 2023 11:56:07 +0200
Subject: [PATCH] Generalized optimizations

---
 compiler/driver.ml                            |  2 +-
 .../{dcalc => shared_ast}/optimizations.ml    | 35 ++++++++++---------
 .../{dcalc => shared_ast}/optimizations.mli   | 16 ++++++---
 compiler/shared_ast/shared_ast.ml             |  1 +
 compiler/verification/conditions.ml           |  6 ++--
 5 files changed, 36 insertions(+), 24 deletions(-)
 rename compiler/{dcalc => shared_ast}/optimizations.ml (90%)
 rename compiler/{dcalc => shared_ast}/optimizations.mli (69%)

diff --git a/compiler/driver.ml b/compiler/driver.ml
index 10d71031..a8b54943 100644
--- a/compiler/driver.ml
+++ b/compiler/driver.ml
@@ -318,7 +318,7 @@ let driver source_file (options : Cli.options) : int =
         let prgm =
           if options.optimize then begin
             Cli.debug_print "Optimizing default calculus...";
-            Dcalc.Optimizations.optimize_program prgm
+            Shared_ast.Optimizations.optimize_program prgm
           end
           else prgm
         in
diff --git a/compiler/dcalc/optimizations.ml b/compiler/shared_ast/optimizations.ml
similarity index 90%
rename from compiler/dcalc/optimizations.ml
rename to compiler/shared_ast/optimizations.ml
index 9d786c8a..e9d5e323 100644
--- a/compiler/dcalc/optimizations.ml
+++ b/compiler/shared_ast/optimizations.ml
@@ -15,21 +15,27 @@
    License for the specific language governing permissions and limitations under
    the License. *)
 open Catala_utils
-open Shared_ast
-open Ast
+open Definitions
 
-type partial_evaluation_ctx = {
-  var_values : (typed expr, typed expr) Var.Map.t;
+type ('a, 'b, 'm) optimizations_ctx = {
+  var_values :
+    ( (('a, 'b) dcalc_lcalc, 'm mark) gexpr,
+      (('a, 'b) dcalc_lcalc, 'm mark) gexpr )
+    Var.Map.t;
   decl_ctx : decl_ctx;
 }
 
-let rec partial_evaluation (ctx : partial_evaluation_ctx) (e : 'm expr) :
-    (dcalc, 'm mark) boxed_gexpr =
+let rec optimize_expr :
+    type a b.
+    (a, b, 'm) optimizations_ctx ->
+    ((a, b) dcalc_lcalc, 'm mark) gexpr ->
+    ((a, b) dcalc_lcalc, 'm mark) boxed_gexpr =
+ fun ctx e ->
   (* We proceed bottom-up, first apply on the subterms *)
-  let e = Expr.map ~f:(partial_evaluation ctx) e in
+  let e = Expr.map ~f:(optimize_expr ctx) e in
   let mark = Marked.get_mark e in
   (* Then reduce the parent node *)
-  let reduce (e : 'm expr) =
+  let reduce (e : ((a, b) dcalc_lcalc, 'm mark) gexpr) =
     (* Todo: improve the handling of eapp(log,elit) cases here, it obfuscates
        the matches and the log calls are not preserved, which would be a good
        property *)
@@ -204,14 +210,11 @@ let rec partial_evaluation (ctx : partial_evaluation_ctx) (e : 'm expr) :
   in
   Expr.Box.app1 e reduce mark
 
-let optimize_expr (decl_ctx : decl_ctx) (e : 'm expr) =
-  partial_evaluation { var_values = Var.Map.empty; decl_ctx } e
+let optimize_expr
+    (decl_ctx : decl_ctx)
+    (e : (('a, 'b) dcalc_lcalc, 'm mark) gexpr) =
+  optimize_expr { var_values = Var.Map.empty; decl_ctx } e
 
 let optimize_program (p : 'm program) : 'm program =
   Bindlib.unbox
-    (Program.map_exprs
-       ~f:
-         (partial_evaluation
-            { var_values = Var.Map.empty; decl_ctx = p.decl_ctx })
-       ~varf:(fun v -> v)
-       p)
+    (Program.map_exprs ~f:(optimize_expr p.decl_ctx) ~varf:(fun v -> v) p)
diff --git a/compiler/dcalc/optimizations.mli b/compiler/shared_ast/optimizations.mli
similarity index 69%
rename from compiler/dcalc/optimizations.mli
rename to compiler/shared_ast/optimizations.mli
index 30491809..8d1a44cf 100644
--- a/compiler/dcalc/optimizations.mli
+++ b/compiler/shared_ast/optimizations.mli
@@ -15,10 +15,16 @@
    License for the specific language governing permissions and limitations under
    the License. *)
 
-(** Optimization passes for default calculus programs and expressions *)
+(** Optimization passes for default calculus and lambda calculus programs and
+    expressions *)
 
-open Shared_ast
-open Ast
+open Definitions
 
-val optimize_expr : decl_ctx -> 'm expr -> (dcalc, 'm mark) boxed_gexpr
-val optimize_program : 'm program -> 'm program
+val optimize_expr :
+  decl_ctx ->
+  (('a, 'b) dcalc_lcalc, 'm mark) gexpr ->
+  (('a, 'b) dcalc_lcalc, 'm mark) boxed_gexpr
+
+val optimize_program :
+  (('a, 'b) dcalc_lcalc, 'm mark) gexpr program ->
+  (('a, 'b) dcalc_lcalc, 'm mark) gexpr program
diff --git a/compiler/shared_ast/shared_ast.ml b/compiler/shared_ast/shared_ast.ml
index 3629b54c..4d825107 100644
--- a/compiler/shared_ast/shared_ast.ml
+++ b/compiler/shared_ast/shared_ast.ml
@@ -24,3 +24,4 @@ module Program = Program
 module Print = Print
 module Typing = Typing
 module Interpreter = Interpreter
+module Optimizations = Optimizations
diff --git a/compiler/verification/conditions.ml b/compiler/verification/conditions.ml
index 7d8de0a1..1c899772 100644
--- a/compiler/verification/conditions.ml
+++ b/compiler/verification/conditions.ml
@@ -351,7 +351,8 @@ let rec generate_verification_conditions_scope_body_expr
         let vc_confl = generate_vc_must_not_return_conflict ctx e in
         let vc_confl =
           if !Cli.optimize_flag then
-            Expr.unbox (Optimizations.optimize_expr ctx.decl vc_confl)
+            Expr.unbox
+              (Shared_ast.Optimizations.optimize_expr ctx.decl vc_confl)
           else vc_confl
         in
         let vc_list =
@@ -373,7 +374,8 @@ let rec generate_verification_conditions_scope_body_expr
             let vc_empty = generate_vc_must_not_return_empty ctx e in
             let vc_empty =
               if !Cli.optimize_flag then
-                Expr.unbox (Optimizations.optimize_expr ctx.decl vc_empty)
+                Expr.unbox
+                  (Shared_ast.Optimizations.optimize_expr ctx.decl vc_empty)
               else vc_empty
             in
             {