From 8a44eff4a5e69513bd3fa2d50d99f332a4ee4ecb Mon Sep 17 00:00:00 2001
From: collin <collin.chin@berkeley.edu>
Date: Tue, 8 Dec 2020 16:00:12 -0500
Subject: [PATCH] fix nested mutable assignee bug, add swap test

---
 compiler/src/statement/assign/assignee.rs |  8 +++++++-
 compiler/tests/mutability/mod.rs          |  8 ++++++++
 compiler/tests/mutability/swap.leo        | 20 ++++++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 compiler/tests/mutability/swap.leo

diff --git a/compiler/src/statement/assign/assignee.rs b/compiler/src/statement/assign/assignee.rs
index 3a0fe2641c..2bb37db9b8 100644
--- a/compiler/src/statement/assign/assignee.rs
+++ b/compiler/src/statement/assign/assignee.rs
@@ -34,7 +34,13 @@ impl<F: Field + PrimeField, G: GroupType<F>> ConstrainedProgram<F, G> {
         // Check that assignee exists and is mutable
         Ok(match self.get_mut(name) {
             Some(value) => match value {
-                ConstrainedValue::Mutable(mutable_value) => mutable_value,
+                ConstrainedValue::Mutable(mutable_value) => {
+                    // Get the mutable value.
+                    mutable_value.get_inner_mut();
+
+                    // Return the mutable value.
+                    mutable_value
+                }
                 _ => return Err(StatementError::immutable_assign(name.to_owned(), span.to_owned())),
             },
             None => return Err(StatementError::undefined_variable(name.to_owned(), span.to_owned())),
diff --git a/compiler/tests/mutability/mod.rs b/compiler/tests/mutability/mod.rs
index 8016bb5412..85d7b1fe8c 100644
--- a/compiler/tests/mutability/mod.rs
+++ b/compiler/tests/mutability/mod.rs
@@ -136,3 +136,11 @@ fn test_function_input_mut() {
 
     assert_satisfied(program);
 }
+
+#[test]
+fn test_swap() {
+    let bytes = include_bytes!("swap.leo");
+    let program = parse_program(bytes).unwrap();
+
+    assert_satisfied(program);
+}
diff --git a/compiler/tests/mutability/swap.leo b/compiler/tests/mutability/swap.leo
new file mode 100644
index 0000000000..8234a3cfb7
--- /dev/null
+++ b/compiler/tests/mutability/swap.leo
@@ -0,0 +1,20 @@
+// Swap two elements of an array.
+function swap(mut a: [u32; 2], i: u32, j: u32) -> [u32; 2] {
+    let t = a[i];
+    a[i] = a[j];
+    a[j] = t;
+    return a
+}
+
+function main() {
+    let mut arr: [u32; 2] = [0, 1];
+    let expected: [u32; 2] = [1, 0];
+
+    // Do swap.
+    let actual = swap(arr, 0, 1);
+
+    // Check result.
+    for i in 0..2 {
+        console.assert(expected[i] == actual[i]);
+    }
+}
\ No newline at end of file