From 6acf49b6feafecf4c468dbba8f9c74f70a118daf Mon Sep 17 00:00:00 2001
From: Nicolas Chataing <nicolas.chataing@ens.fr>
Date: Sun, 9 Aug 2020 23:50:00 +0200
Subject: [PATCH] Rewrite merge_var_redefs function

---
 src/catala/ir/lambda.ml               |  4 ++++
 src/catala/translation/interpreter.ml | 34 +++++++++++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/catala/ir/lambda.ml b/src/catala/ir/lambda.ml
index 99acf84e..9b6f912a 100644
--- a/src/catala/ir/lambda.ml
+++ b/src/catala/ir/lambda.ml
@@ -61,6 +61,10 @@ let get_typ ((_, typ) : term) : typ = typ
 let map_untype (f : untyped_term -> untyped_term) (((term, pos), typ) : term) : term =
   ((f term, pos), typ)
 
+let map_untype2 (f : untyped_term -> untyped_term -> untyped_term) (((t1, pos), typ) : term)
+    (((t2, _), _) : term) : term =
+  ((f t1 t2, pos), typ)
+
 let empty_default_term : default_term = { defaults = IntMap.empty; ordering = []; nb_defaults = 0 }
 
 let add_default (just : term) (cons : term) (term : default_term) =
diff --git a/src/catala/translation/interpreter.ml b/src/catala/translation/interpreter.ml
index 67b9c30e..f1720dda 100644
--- a/src/catala/translation/interpreter.ml
+++ b/src/catala/translation/interpreter.ml
@@ -214,8 +214,38 @@ let build_scope_schedule (ctxt : Context.context) (scope : Scope.scope) : G.t =
     scope.scope_sub_defs;
   g
 
-let merge_var_redefs (_subscope : Scope.scope) (_redefs : Scope.definition UidMap.t) : Scope.scope =
-  assert false
+let merge_var_redefs (subscope : Scope.scope) (redefs : Scope.definition UidMap.t) : Scope.scope =
+  let merge_defaults : Lambda.term -> Lambda.term -> Lambda.term =
+    Lambda.map_untype2 (fun old_t new_t ->
+        match (old_t, new_t) with
+        | EDefault old_def, EDefault new_def ->
+            EDefault (Lambda.merge_default_terms old_def new_def)
+        | EFun ([ bind ], old_t), EFun (_, new_t) ->
+            let body =
+              Lambda.map_untype2
+                (fun old_t new_t ->
+                  match (old_t, new_t) with
+                  | EDefault old_def, EDefault new_def ->
+                      EDefault (Lambda.merge_default_terms old_def new_def)
+                  | _ -> assert false)
+                old_t new_t
+            in
+            EFun ([ bind ], body)
+        | _ -> assert false)
+  in
+
+  {
+    subscope with
+    scope_defs =
+      UidMap.fold
+        (fun uid new_def sub_defs ->
+          match UidMap.find_opt uid sub_defs with
+          | None -> UidMap.add uid new_def sub_defs
+          | Some old_def ->
+              let def = merge_defaults old_def new_def in
+              UidMap.add uid def sub_defs)
+        redefs subscope.scope_defs;
+  }
 
 (*{ subscope with scope_defs = UidMap.fold (fun uid new_def sub_defs -> match UidMap.find_opt uid
   sub_defs with | None -> UidMap.add uid new_def sub_defs | Some old_def -> let def =