Fix a bug in closure conversion

compiler/driver.ml

@@ -435,6 +435,18 @@ let driver source_file (options : Cli.options) : int =
                 Messages.emit_debug "Performing closure conversion...";
                 let prgm = Lcalc.Closure_conversion.closure_conversion prgm in
                 let prgm = Bindlib.unbox prgm in
+                let _output_file, with_output = get_output_format () in
+                with_output
+                @@ fun fmt ->
+                if Option.is_some options.ex_scope then
+                  Format.fprintf fmt "%a\n"
+                    (Shared_ast.Print.scope ~debug:options.debug prgm.decl_ctx)
+                    (scope_uid, Shared_ast.Program.get_scope_body prgm scope_uid)
+                else
+                  Format.fprintf fmt "%a\n"
+                    (Shared_ast.Print.program ~debug:options.debug)
+                    prgm;
+                Format.pp_print_flush Format.std_formatter ();
                 let prgm =
                   if options.optimize then (
                     Messages.emit_debug "Optimizing lambda calculus...";

compiler/lcalc/closure_conversion.ml

@@ -37,12 +37,9 @@ let rec hoist_context_free_closures :
  fun ctx e ->
   let m = Mark.get e in
   match Mark.remove e with
-  | EStruct _ | EStructAccess _ | ETuple _ | ETupleAccess _ | EInj _ | EArray _
-  | ELit _ | EAssert _ | EOp _ | EIfThenElse _ | ERaise _ | ECatch _ | EVar _ ->
-    Expr.map_gather ~acc:[] ~join:( @ ) ~f:(hoist_context_free_closures ctx) e
   | EMatch { e; cases; name } ->
     let collected_closures, new_e = (hoist_context_free_closures ctx) e in
-    (* We do not close the clotures inside the arms of the match expression,
+    (* We do not close the closures inside the arms of the match expression,
        since they get a special treatment at compilation to Scalc. *)
     let collected_closures, new_cases =
       EnumConstructor.Map.fold
@@ -83,7 +80,9 @@ let rec hoist_context_free_closures :
     (* this is the closure we want to hoist*)
     let closure_var = Var.make ctx.name_context in
     [{ name = closure_var; closure = e }], Expr.make_var closure_var m
-  | EApp _ ->
+  | EApp _ | EStruct _ | EStructAccess _ | ETuple _ | ETupleAccess _ | EInj _
+  | EArray _ | ELit _ | EAssert _ | EOp _ | EIfThenElse _ | ERaise _ | ECatch _
+  | EVar _ ->
     Expr.map_gather ~acc:[] ~join:( @ ) ~f:(hoist_context_free_closures ctx) e
   | _ -> .
  [@@warning "-32"]
@@ -118,7 +117,9 @@ let rec transform_closures_expr :
             let vars, body = Bindlib.unmbind binder in
             let new_free_vars, new_body = (transform_closures_expr ctx) body in
             let new_binder = Expr.bind vars new_body in
-            ( Var.Set.union free_vars new_free_vars,
+            ( Var.Set.union free_vars
+                (Var.Set.diff new_free_vars
+                   (Var.Set.of_list (Array.to_list vars))),
               EnumConstructor.Map.add cons
                 (Expr.eabs new_binder tys (Mark.get e1))
                 new_cases )

compiler/shared_ast/print.ml

@@ -546,7 +546,7 @@ let rec expr_aux :
   | EApp { f; args } ->
     Format.fprintf fmt "@[<hv 2>%a@ %a@]" (lhs exprc) f
       (Format.pp_print_list
-         ~pp_sep:(fun fmt () -> Format.fprintf fmt ",@ ")
+         ~pp_sep:(fun fmt () -> Format.fprintf fmt "@ ")
          (rhs exprc))
       args
   | EIfThenElse _ ->

tests/test_array/good/aggregation_3.catala_en

@@ -41,27 +41,27 @@ let scope S (x: integer|internal|output) =
             filter (λ (i: integer) → i > 2) [ 1; 2; 3 ])
          = [ 5 ];
   assert (reduce
-            (λ (sum1: integer) (sum2: integer) → sum1 + sum2),
-            0,
+            (λ (sum1: integer) (sum2: integer) → sum1 + sum2)
+            0
             [ 1; 2; 3 ])
          = 6;
   assert (reduce
-            (λ (sum1: integer) (sum2: integer) → sum1 + sum2),
-            0,
+            (λ (sum1: integer) (sum2: integer) → sum1 + sum2)
+            0
             map (λ (i: integer) → i + 2) [ 1; 2; 3 ])
          = 12;
   assert (length [ 1; 2; 3 ]) = 3;
   assert (length filter (λ (i: integer) → i >= 2) [ 1; 2; 3 ]) = 2;
   assert (reduce
             (λ (max1: integer) (max2: integer) →
-             if max1 > max2 then max1 else max2),
-            10,
+             if max1 > max2 then max1 else max2)
+            10
             [ 1; 2; 3 ])
          = 3;
   assert (reduce
             (λ (max1: decimal) (max2: decimal) →
-             if max1 > max2 then max1 else max2),
-            10.,
+             if max1 > max2 then max1 else max2)
+            10.
             map (λ (i: integer) → to_rat i) [ 1; 2; 3 ])
          = 3.;
   assert (reduce
@@ -72,8 +72,8 @@ let scope S (x: integer|internal|output) =
                < let i : integer = i_2 in
                  to_rat ((2 - i) * (2 - i))
              then i_1
-             else i_2),
-            42,
+             else i_2)
+            42
             [ 1; 2; 3 ])
          = 2
 ```

tests/test_func/good/closure_conversion.catala_en

@@ -13,12 +13,61 @@ scope S:
 
 ```catala-test-inline
 $ catala Lcalc -s S --avoid_exceptions -O --closure_conversion
-[ERROR] Variable x not found in the current context
-
-┌─⯈ tests/test_func/good/closure_conversion.catala_en:5.12-5.13:
-└─┐
-5 │   internal f content integer depends on y content integer
-  │            ‾
-  └─ Article
-#return code 255#
+let scope S (S_in: S_in {x_in: option bool}): S {z: option integer} =
+  let get x : any = S_in.x_in in
+  let set f : any =
+    ESome
+      let f : any =
+        λ (env: any) (y: integer) →
+        let x : any = env.0 in
+        ESome
+          match
+            (match x with
+             | ENone _ → ENone _
+             | ESome x1 → if x1 then ESome y else ESome - y)
+            with
+          | ENone _ → raise NoValueProvided
+          | ESome f → f
+      in
+      (f, (x))
+  in
+  let set z : any =
+    ESome
+      match
+        (match f with
+         | ENone _ → ENone _
+         | ESome f →
+           let code_and_env : any = f in
+           let code : any = code_and_env.0 in
+           let env : any = code_and_env.1 in
+           code env -1)
+        with
+      | ENone _ → raise NoValueProvided
+      | ESome z → z
+  in
+  return { S z = z; }
+let scope S (S_in: S_in {x_in: option bool}): S {z: option integer} =
+  let get x : option bool = S_in.x_in in
+  let set f :
+      option (((option bool), integer) → option integer * (option bool)) =
+    ESome
+      (λ (env: (option bool)) (y: integer) →
+        ESome
+          match
+            (match env.0 with
+             | ENone _ → ENone _
+             | ESome x → if x then ESome y else ESome - y)
+            with
+          | ENone _ → raise NoValueProvided
+          | ESome f → f, (x))
+  in
+  let set z : option integer =
+    ESome
+      match (match f with
+             | ENone _ → ENone _
+             | ESome f → f.0 f.1 -1) with
+      | ENone _ → raise NoValueProvided
+      | ESome z → z
+  in
+  return { S z = z; }
 ```

tests/test_scope/good/scope_call4.catala_en

@@ -46,10 +46,10 @@ f1 =
        match
          (match (ESome (λ (x1: integer) → ESome (x1 + 1))) with
           | ENone _ → ENone _
-          | ESome f1 → f1 (x + 1))
+          | ESome g → g (x + 1))
          with
-       | ENone f1 → raise NoValueProvided
-       | ESome x1 → x1)
+       | ENone _ → raise NoValueProvided
+       | ESome f1 → f1)
 [RESULT]
 f2 =
   ESome
@@ -58,8 +58,8 @@ f2 =
        match
          (match (ESome (λ (x1: integer) → ESome (x1 + 1))) with
           | ENone _ → ENone _
-          | ESome f2 → f2 (x + 1))
+          | ESome g → g (x + 1))
          with
-       | ENone f2 → raise NoValueProvided
-       | ESome x1 → x1)
+       | ENone _ → raise NoValueProvided
+       | ESome f2 → f2)
 ```

tests/test_scope/good/simple.catala_en

@@ -12,7 +12,7 @@ scope Foo:
 $ catala Lcalc -s Foo
 let scope Foo (Foo_in: Foo_in): Foo {bar: integer} =
   let set bar : integer =
-    try handle_default [  ], (λ (_: unit) → true), (λ (_: unit) → 0)
+    try handle_default [  ] (λ (_: unit) → true) (λ (_: unit) → 0)
     with EmptyError -> raise NoValueProvided
   in
   return { Foo bar = bar; }