From a56db3ec9b20de587735e2f002be5c355c6b6b83 Mon Sep 17 00:00:00 2001
From: Alex Vitkov <44268717+alexvitkov@users.noreply.github.com>
Date: Mon, 7 Aug 2023 19:04:31 +0300
Subject: [PATCH] fix: Fix an ICE when reassigning a mutable lambda variable to
 one with a different environment type (#2172)

* fix: don't allow lambda reassignment if closure types dont match

* fix: add test for bad lambda reassignment

* fix: cargo fmt & clippy

* fix: alternate approach

---------

Co-authored-by: Alex Vitkov <alex.vitkov@metacraft-labs.com>
---
 .../fail/different-lambda-env-reassign-disallow.nr       | 9 +++++++++
 crates/noirc_frontend/src/hir_def/types.rs               | 8 ++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)
 create mode 100644 crates/nargo_cli/tests/compile_tests_data/fail/different-lambda-env-reassign-disallow.nr

diff --git a/crates/nargo_cli/tests/compile_tests_data/fail/different-lambda-env-reassign-disallow.nr b/crates/nargo_cli/tests/compile_tests_data/fail/different-lambda-env-reassign-disallow.nr
new file mode 100644
index 00000000000..7bac2367846
--- /dev/null
+++ b/crates/nargo_cli/tests/compile_tests_data/fail/different-lambda-env-reassign-disallow.nr
@@ -0,0 +1,9 @@
+fn bad()  {
+    let a: i32 = 100;
+    let b: i32 = 200;
+
+    let mut f = || a;
+
+    // this should fail with a type error, since the closures have different environments & types
+    f = || a + b; 
+}
\ No newline at end of file
diff --git a/crates/noirc_frontend/src/hir_def/types.rs b/crates/noirc_frontend/src/hir_def/types.rs
index d77b8033ba1..587001bfafc 100644
--- a/crates/noirc_frontend/src/hir_def/types.rs
+++ b/crates/noirc_frontend/src/hir_def/types.rs
@@ -1206,12 +1206,14 @@ impl Type {
                 }
             }
 
-            (Function(params_a, ret_a, _env_a), Function(params_b, ret_b, _env_b)) => {
+            (Function(params_a, ret_a, env_a), Function(params_b, ret_b, env_b)) => {
                 if params_a.len() == params_b.len() {
                     for (a, b) in params_a.iter().zip(params_b.iter()) {
                         a.try_unify(b, span)?;
                     }
 
+                    env_a.try_unify(env_b, span)?;
+
                     ret_b.try_unify(ret_a, span)
                 } else {
                     Err(SpanKind::None)
@@ -1413,12 +1415,14 @@ impl Type {
                 }
             }
 
-            (Function(params_a, ret_a, _env_a), Function(params_b, ret_b, _env_b)) => {
+            (Function(params_a, ret_a, env_a), Function(params_b, ret_b, env_b)) => {
                 if params_a.len() == params_b.len() {
                     for (a, b) in params_a.iter().zip(params_b) {
                         a.is_subtype_of(b, span)?;
                     }
 
+                    env_a.is_subtype_of(env_b, span)?;
+
                     // return types are contravariant, so this must be ret_b <: ret_a instead of the reverse
                     ret_b.is_subtype_of(ret_a, span)
                 } else {