corrected iota optimization

compiler/dcalc/invariant.ml

@@ -55,6 +55,23 @@ let invariant_app_is_either_op_var_let () : string * invariant_expr =
       | EApp _ -> Fail
       | _ -> Ignore )
 
+(** the arity of constructors when matching is always one. *)
+let invariant_match () : string * invariant_expr =
+  ( __FUNCTION__,
+    fun e ->
+      match Marked.unmark e with
+      | EMatch { cases; _ } ->
+        if
+          EnumConstructor.Map.for_all
+            (fun _ case ->
+              match Marked.unmark case with
+              | EAbs { binder; _ } -> Bindlib.mbinder_arity binder = 1
+              | _ -> false)
+            cases
+        then Pass
+        else Fail
+      | _ -> Ignore )
+
 let check_invariant (inv : string * invariant_expr) (p : typed program) : bool =
   (* TODO: add a Program.fold_exprs to get rid of the reference 0:-)? *)
   let result = ref true in

compiler/driver.ml

@@ -284,6 +284,7 @@ let driver source_file (options : Cli.options) : int =
                   check_invariant (invariant_no_partial_evaluation ()) prgm;
                   check_invariant (invariant_no_return_a_function ()) prgm;
                   check_invariant (invariant_app_is_either_op_var_let ()) prgm;
+                  check_invariant (invariant_match ()) prgm;
                 ]
             in
 

compiler/lcalc/optimizations.ml

@@ -31,12 +31,28 @@ let rec iota_expr (e : 'm expr) : 'm expr boxed =
     Expr.eapp case [e1] m
   | EMatch { e = e'; cases; name = n }
     when cases
-         |> EnumConstructor.Map.mapi (fun i case ->
+         |> EnumConstructor.Map.for_all (fun i case ->
                 match Marked.unmark case with
-                | EInj { cons = i'; name = n'; _ } ->
-                  EnumConstructor.equal i i' && EnumName.equal n n'
-                | _ -> false)
-         |> EnumConstructor.Map.for_all (fun _ b -> b) ->
+                | EAbs { binder; _ } -> (
+                  let var, body = Bindlib.unmbind binder in
+                  (* because of invariant [invariant_match], the arity is always
+                     one. *)
+                  let[@warning "-8"] [| var |] = var in
+
+                  match Marked.unmark body with
+                  | EInj { cons = i'; name = n'; e = EVar x, _ } ->
+                    EnumConstructor.equal i i'
+                    && EnumName.equal n n'
+                    && Bindlib.eq_vars x var
+                  | EInj { cons = i'; name = n'; e = ELit LUnit, _ } ->
+                    (* since unit is the only value of type unit. We don't need
+                       to check the equality. *)
+                    EnumConstructor.equal i i' && EnumName.equal n n'
+                  | _ -> false)
+                | _ ->
+                  (* because of invariant [invariant_match], there is always
+                     some EAbs in each cases. *)
+                  assert false) ->
     visitor_map iota_expr e'
   | _ -> visitor_map iota_expr e