From a310f3c245144b235815cf59c7a4e125294ca005 Mon Sep 17 00:00:00 2001
From: lynskylate <lynskylate@gmail.com>
Date: Wed, 8 Jul 2020 01:06:12 +0800
Subject: [PATCH] Add annotation setup

---
 bytecode/src/bytecode.rs |  2 ++
 compiler/src/compile.rs  | 37 ++++++++++++++-----------------------
 vm/src/frame.rs          |  7 +++++++
 3 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
index 835905c1b..09d0e03d5 100644
--- a/bytecode/src/bytecode.rs
+++ b/bytecode/src/bytecode.rs
@@ -205,6 +205,7 @@ pub enum Instruction {
     ReturnValue,
     YieldValue,
     YieldFrom,
+    SetupAnnotation,
     SetupLoop {
         start: Label,
         end: Label,
@@ -570,6 +571,7 @@ impl Instruction {
             ReturnValue => w!(ReturnValue),
             YieldValue => w!(YieldValue),
             YieldFrom => w!(YieldFrom),
+            SetupAnnotation => w!(SetupAnnotation),
             SetupLoop { start, end } => w!(SetupLoop, label_map[start], label_map[end]),
             SetupExcept { handler } => w!(SetupExcept, label_map[handler]),
             SetupFinally { handler } => w!(SetupFinally, label_map[handler]),
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
index af3d676dc..ee81ff7db 100644
--- a/compiler/src/compile.rs
+++ b/compiler/src/compile.rs
@@ -984,15 +984,14 @@ impl<O: OutputStream> Compiler<O> {
         Ok(())
     }
 
-    fn option_stmt_to_bool(&mut self, stmts: &Option<ast::Suite>) -> bool {
-        match &stmts {
-            Some(stmts) => self.find_ann(stmts),
-            None => false,
-        }
-    }
-
-    fn find_ann(&mut self, body: &[ast::Statement]) -> bool {
+    fn find_ann(&self, body: &[ast::Statement]) -> bool {
         use ast::StatementType::*;
+        let option_stmt_to_bool = |suit: &Option<ast::Suite>| -> bool {
+            match suit {
+                Some(stmts) => self.find_ann(stmts),
+                None => false,
+            }
+        };
 
         for statement in body {
             let res = match &statement.node {
@@ -1007,17 +1006,17 @@ impl<O: OutputStream> Compiler<O> {
                     iter: _,
                     body,
                     orelse,
-                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                } => self.find_ann(body) || option_stmt_to_bool(orelse),
                 If {
                     test: _,
                     body,
                     orelse,
-                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                } => self.find_ann(body) || option_stmt_to_bool(orelse),
                 While {
                     test: _,
                     body,
                     orelse,
-                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                } => self.find_ann(body) || option_stmt_to_bool(orelse),
                 With {
                     is_async: _,
                     items: _,
@@ -1029,9 +1028,9 @@ impl<O: OutputStream> Compiler<O> {
                     orelse,
                     finalbody,
                 } => {
-                    self.find_ann(body)
-                        || self.option_stmt_to_bool(orelse)
-                        || self.option_stmt_to_bool(finalbody)
+                    self.find_ann(&body)
+                        || option_stmt_to_bool(orelse)
+                        || option_stmt_to_bool(finalbody)
                 }
                 _ => false,
             };
@@ -1097,15 +1096,7 @@ impl<O: OutputStream> Compiler<O> {
         });
         // setup annotations
         if self.find_ann(body) {
-            self.emit(Instruction::BuildMap {
-                size: 0,
-                unpack: false,
-                for_call: false,
-            });
-            self.emit(Instruction::StoreName {
-                name: "__annotations__".to_owned(),
-                scope: bytecode::NameScope::Free,
-            });
+            self.emit(Instruction::SetupAnnotation);
         }
         self.compile_statements(new_body)?;
         self.emit(Instruction::LoadConst {
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index eef0147ce..8b8122b0a 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -456,6 +456,13 @@ impl ExecutingFrame<'_> {
                 Ok(Some(ExecutionResult::Yield(value)))
             }
             bytecode::Instruction::YieldFrom => self.execute_yield_from(vm),
+            bytecode::Instruction::SetupAnnotation => {
+                let locals = self.scope.get_locals();
+                if let Err(_e) = locals.get_item("__annotations__", vm) {
+                    locals.set_item("__annotations__", vm.ctx.new_dict().into_object(), vm)?;
+                }
+                Ok(None)
+            }
             bytecode::Instruction::SetupLoop { start, end } => {
                 self.push_block(BlockType::Loop {
                     start: *start,