From 928b3ad5d93943960cc6f480b28bce25f29b3271 Mon Sep 17 00:00:00 2001
From: Alexander Ivanov <alehander42@gmail.com>
Date: Fri, 4 Aug 2023 18:11:01 +0300
Subject: [PATCH] fix: fix an ICE happening when we call a closure result from
 if/else (#2146)

* fix: fix an ICE happening when we call a closure result from if/else

* refactor: moved test to a regression fn, cleanup code
---
 .../higher_order_functions/src/main.nr        | 20 +++++++++++++++++++
 .../src/monomorphization/mod.rs               | 15 ++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/crates/nargo_cli/tests/test_data/higher_order_functions/src/main.nr b/crates/nargo_cli/tests/test_data/higher_order_functions/src/main.nr
index fefd23b7dbc..782b6af998e 100644
--- a/crates/nargo_cli/tests/test_data/higher_order_functions/src/main.nr
+++ b/crates/nargo_cli/tests/test_data/higher_order_functions/src/main.nr
@@ -36,6 +36,8 @@ fn main() -> pub Field {
     x += 1;
     assert(closure_capturing_mutable(1) == 5);
 
+    regression_2154();
+
     let ret = twice(add1, 3);
 
     test_array_functions();
@@ -85,3 +87,21 @@ fn add1(x: Field) -> Field {
 fn twice(f: fn(Field) -> Field, x: Field) -> Field {
     f(f(x))
 }
+
+// Fixing an ICE, where rewriting the closures
+// during monomorphization didn't correspond
+// to an internal `if` type
+// found by @jfecher:
+// https://github.com/noir-lang/noir/pull/1959#issuecomment-1658992989
+// issue https://github.com/noir-lang/noir/issues/2154
+fn regression_2154() {
+    let x: u32 = 32;
+
+    let closure_if_else = if x > 2 {
+        || x
+    } else {
+        || x + 2342
+    };
+
+    assert(closure_if_else() == 32);
+}
diff --git a/crates/noirc_frontend/src/monomorphization/mod.rs b/crates/noirc_frontend/src/monomorphization/mod.rs
index c8167baf6bb..6c9724de6ce 100644
--- a/crates/noirc_frontend/src/monomorphization/mod.rs
+++ b/crates/noirc_frontend/src/monomorphization/mod.rs
@@ -689,8 +689,19 @@ impl<'interner> Monomorphizer<'interner> {
             HirType::Function(args, ret, env) => {
                 let args = vecmap(args, Self::convert_type);
                 let ret = Box::new(Self::convert_type(ret));
-                let env = Box::new(Self::convert_type(env));
-                ast::Type::Function(args, ret, env)
+                let env = Self::convert_type(env);
+                match &env {
+                    ast::Type::Unit => ast::Type::Function(args, ret, Box::new(env)),
+                    ast::Type::Tuple(_elements) => ast::Type::Tuple(vec![
+                        env.clone(),
+                        ast::Type::Function(args, ret, Box::new(env)),
+                    ]),
+                    _ => {
+                        unreachable!(
+                            "internal Type::Function env should be either a Unit or a Tuple, not {env}"
+                        )
+                    }
+                }
             }
 
             HirType::MutableReference(element) => {