Generalized optimizations

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

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)

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

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

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
             {