From 6abf1511e99e0eeb543445f0b7d88eb57ea8a7f5 Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Tue, 16 Apr 2019 17:19:57 +0200
Subject: [PATCH 1/2] Add nonlocal support.

---
 tests/snippets/global_nonlocal.py | 10 ++++++++++
 vm/src/bytecode.rs                |  1 +
 vm/src/compile.rs                 |  2 +-
 vm/src/frame.rs                   | 17 ++++++++++++++++-
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/tests/snippets/global_nonlocal.py b/tests/snippets/global_nonlocal.py
index 48d1242108..fd5a61e1cc 100644
--- a/tests/snippets/global_nonlocal.py
+++ b/tests/snippets/global_nonlocal.py
@@ -11,3 +11,13 @@ def b():
 b()
 assert a == 4
 
+def x():
+    def y():
+        nonlocal b
+        b = 3
+    b = 2
+    y()
+    return b
+
+res = x()
+assert res == 3, str(res)
diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs
index 10c667890e..73d69dc180 100644
--- a/vm/src/bytecode.rs
+++ b/vm/src/bytecode.rs
@@ -41,6 +41,7 @@ pub type Label = usize;
 #[derive(Debug, Clone, PartialEq)]
 pub enum NameScope {
     Local,
+    NonLocal,
     Global,
 }
 
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
index d28c10feb4..8635e21ac8 100644
--- a/vm/src/compile.rs
+++ b/vm/src/compile.rs
@@ -179,12 +179,12 @@ impl Compiler {
         let role = self.lookup_name(name);
         match role {
             SymbolRole::Global => bytecode::NameScope::Global,
+            SymbolRole::Nonlocal => bytecode::NameScope::NonLocal,
             _ => bytecode::NameScope::Local,
         }
     }
 
     fn load_name(&mut self, name: &str) {
-        // TODO: if global, do something else!
         let scope = self.scope_for_name(name);
         self.emit(Instruction::LoadName {
             name: name.to_string(),
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index 2bf8215c1c..4425a8a7d5 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -129,6 +129,7 @@ pub trait NameProtocol {
     fn store_name(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
     fn delete_name(&self, vm: &VirtualMachine, name: &str);
     fn load_cell(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
+    fn store_cell(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
     fn load_global(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
     fn store_global(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
 }
@@ -157,6 +158,16 @@ impl NameProtocol for Scope {
         None
     }
 
+    fn store_cell(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef) {
+        self.locals
+            .iter()
+            .skip(1)
+            .next()
+            .unwrap()
+            .set_item(name, value, vm)
+            .unwrap();
+    }
+
     fn store_name(&self, vm: &VirtualMachine, key: &str, value: PyObjectRef) {
         self.get_locals().set_item(key, value, vm).unwrap();
     }
@@ -1037,8 +1048,11 @@ impl Frame {
             bytecode::NameScope::Global => {
                 self.scope.store_global(vm, name, obj);
             }
+            bytecode::NameScope::NonLocal => {
+                self.scope.store_cell(vm, name, obj);
+            }
             bytecode::NameScope::Local => {
-                self.scope.store_name(&vm, name, obj);
+                self.scope.store_name(vm, name, obj);
             }
         }
         Ok(None)
@@ -1057,6 +1071,7 @@ impl Frame {
     ) -> FrameResult {
         let optional_value = match name_scope {
             bytecode::NameScope::Global => self.scope.load_global(vm, name),
+            bytecode::NameScope::NonLocal => self.scope.load_cell(vm, name),
             bytecode::NameScope::Local => self.scope.load_name(&vm, name),
         };
 

From 6d1b807c77a5c1d1e5e7afc78729cde9a6cfdef6 Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Wed, 17 Apr 2019 12:59:56 +0200
Subject: [PATCH 2/2] Improve error message of unwrap operation on nonlocal
 scope.

---
 vm/src/frame.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index 4425a8a7d5..9d3e28e1d5 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -163,7 +163,7 @@ impl NameProtocol for Scope {
             .iter()
             .skip(1)
             .next()
-            .unwrap()
+            .expect("no outer scope for non-local")
             .set_item(name, value, vm)
             .unwrap();
     }