From fae42a053309b90388f5b61bd514a40371080122 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

---
 src/compile.rs | 37 ++++++++++++++-----------------------
 1 file changed, 14 insertions(+), 23 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index af3d676..ee81ff7 100644
--- a/src/compile.rs
+++ b/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 {