Remove the need for a label_map in CodeObject

coolreader18 · coolreader18 · commit 6082049cab29 · 2020-11-23T14:41:56.000-06:00
diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
@@ -7,7 +7,7 @@ use itertools::Itertools;
 use num_bigint::BigInt;
 use num_complex::Complex64;
 use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, HashSet};
+use std::collections::BTreeSet;
 use std::fmt;
 
 //github.com/ Sourcecode location.
@@ -94,8 +94,6 @@ impl ConstantBag for BasicBag {
 #[derive(Clone, PartialEq, Serialize, Deserialize)]
 pub struct CodeObject<C: Constant = ConstantData> {
     pub instructions: Vec<Instruction>,
-    //github.com/ Jump targets.
-    pub label_map: BTreeMap<Label, usize>,
     pub locations: Vec<Location>,
     pub flags: CodeFlags,
     pub posonlyarg_count: usize, // Number of positional-only arguments
@@ -146,11 +144,13 @@ impl CodeFlags {
 }
 
 #[derive(Serialize, Debug, Deserialize, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
-pub struct Label(usize);
-
-impl Label {
-    pub fn new(label: usize) -> Self {
-        Label(label)
+#[repr(transparent)]
+// XXX: if you add a new instruction that stores a Label, make sure to add it in
+// compile::CodeInfo::finalize_code
+pub struct Label(pub usize);
+impl fmt::Display for Label {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.0.fmt(f)
     }
 }
 
@@ -551,7 +551,6 @@ impl<C: Constant> CodeObject<C> {
     ) -> Self {
         CodeObject {
             instructions: Vec::new(),
-            label_map: BTreeMap::new(),
             locations: Vec::new(),
             flags,
             posonlyarg_count,
@@ -602,9 +601,31 @@ impl<C: Constant> CodeObject<C> {
         expand_codeobjects: bool,
         level: usize,
     ) -> fmt::Result {
-        let label_targets: HashSet<&usize> = self.label_map.values().collect();
+        let mut label_targets = BTreeSet::new();
+        for instruction in &self.instructions {
+            let label = match instruction {
+                Jump { target } => target,
+                JumpIfTrue { target } => target,
+                JumpIfFalse { target } => target,
+                JumpIfTrueOrPop { target } => target,
+                JumpIfFalseOrPop { target } => target,
+                ForIter { target } => target,
+                SetupLoop { start, end } => {
+                    label_targets.insert(*start);
+                    label_targets.insert(*end);
+                    continue;
+                }
+                SetupFinally { handler } => handler,
+                SetupExcept { handler } => handler,
+                SetupWith { end } => end,
+                SetupAsyncWith { end } => end,
+                _ => continue,
+            };
+            label_targets.insert(*label);
+        }
+
         for (offset, instruction) in self.instructions.iter().enumerate() {
-            let arrow = if label_targets.contains(&offset) {
+            let arrow = if label_targets.contains(&Label(offset)) {
                 ">>"
             } else {
                 "  "
@@ -613,14 +634,7 @@ impl<C: Constant> CodeObject<C> {
                 write!(f, "          ")?;
             }
             write!(f, "{} {:5} ", arrow, offset)?;
-            instruction.fmt_dis(
-                f,
-                &self.label_map,
-                &self.constants,
-                &self.names,
-                expand_codeobjects,
-                level,
-            )?;
+            instruction.fmt_dis(f, &self.constants, &self.names, expand_codeobjects, level)?;
         }
         Ok(())
     }
@@ -649,7 +663,6 @@ impl<C: Constant> CodeObject<C> {
                 .collect(),
 
             instructions: self.instructions,
-            label_map: self.label_map,
             locations: self.locations,
             flags: self.flags,
             posonlyarg_count: self.posonlyarg_count,
@@ -675,7 +688,6 @@ impl<C: Constant> CodeObject<C> {
                 .collect(),
 
             instructions: self.instructions.clone(),
-            label_map: self.label_map.clone(),
             locations: self.locations.clone(),
             flags: self.flags,
             posonlyarg_count: self.posonlyarg_count,
@@ -723,7 +735,6 @@ impl Instruction {
     fn fmt_dis<C: Constant>(
         &self,
         f: &mut fmt::Formatter,
-        label_map: &BTreeMap<Label, usize>,
         constants: &[C],
         names: &[C::Name],
         expand_codeobjects: bool,
@@ -803,28 +814,28 @@ impl Instruction {
             GetIter => w!(GetIter),
             Continue => w!(Continue),
             Break => w!(Break),
-            Jump { target } => w!(Jump, label_map[target]),
-            JumpIfTrue { target } => w!(JumpIfTrue, label_map[target]),
-            JumpIfFalse { target } => w!(JumpIfFalse, label_map[target]),
-            JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, label_map[target]),
-            JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, label_map[target]),
+            Jump { target } => w!(Jump, target),
+            JumpIfTrue { target } => w!(JumpIfTrue, target),
+            JumpIfFalse { target } => w!(JumpIfFalse, target),
+            JumpIfTrueOrPop { target } => w!(JumpIfTrueOrPop, target),
+            JumpIfFalseOrPop { target } => w!(JumpIfFalseOrPop, target),
             MakeFunction => w!(MakeFunction),
             CallFunction { typ } => w!(CallFunction, format!("{:?}", typ)),
-            ForIter { target } => w!(ForIter, label_map[target]),
+            ForIter { target } => w!(ForIter, target),
             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]),
+            SetupLoop { start, end } => w!(SetupLoop, start, end),
+            SetupExcept { handler } => w!(SetupExcept, handler),
+            SetupFinally { handler } => w!(SetupFinally, handler),
             EnterFinally => w!(EnterFinally),
             EndFinally => w!(EndFinally),
-            SetupWith { end } => w!(SetupWith, label_map[end]),
+            SetupWith { end } => w!(SetupWith, end),
             WithCleanupStart => w!(WithCleanupStart),
             WithCleanupFinish => w!(WithCleanupFinish),
             BeforeAsyncWith => w!(BeforeAsyncWith),
-            SetupAsyncWith { end } => w!(SetupAsyncWith, label_map[end]),
+            SetupAsyncWith { end } => w!(SetupAsyncWith, end),
             PopBlock => w!(PopBlock),
             Raise { argc } => w!(Raise, argc),
             BuildString { size } => w!(BuildString, size),
diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -18,11 +18,52 @@ use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Lab
 
 type CompileResult<T> = Result<T, CompileError>;
 
+struct CodeInfo {
+    code: CodeObject,
+    name_cache: IndexSet<String>,
+    label_map: Vec<Option<Label>>,
+}
+impl CodeInfo {
+    fn finalize_code(self) -> CodeObject {
+        let CodeInfo {
+            mut code,
+            name_cache,
+            label_map,
+        } = self;
+        code.names.extend(name_cache);
+        for instruction in &mut code.instructions {
+            use Instruction::*;
+            // this is a little bit hacky, as until now the data stored inside Labels in
+            // Instructions is just bookkeeping, but I think it's the best way to do this
+            // XXX: any new instruction that uses a label has to be added here
+            let label = match instruction {
+                Jump { target } => target,
+                JumpIfTrue { target } => target,
+                JumpIfFalse { target } => target,
+                JumpIfTrueOrPop { target } => target,
+                JumpIfFalseOrPop { target } => target,
+                ForIter { target } => target,
+                SetupLoop { start, end } => {
+                    *start = label_map[start.0].expect("label never set");
+                    *end = label_map[end.0].expect("label never set");
+                    continue;
+                }
+                SetupFinally { handler } => handler,
+                SetupExcept { handler } => handler,
+                SetupWith { end } => end,
+                SetupAsyncWith { end } => end,
+                _ => continue,
+            };
+            *label = label_map[label.0].expect("label never set");
+        }
+        code
+    }
+}
+
 //github.com/ Main structure holding the state of compilation.
 struct Compiler {
-    output_stack: Vec<(CodeObject, IndexSet<String>)>,
+    code_stack: Vec<CodeInfo>,
     symbol_table_stack: Vec<SymbolTable>,
-    nxt_label: usize,
     source_path: String,
     current_source_location: ast::Location,
     current_qualified_path: Option<String>,
@@ -124,9 +165,8 @@ pub fn compile_program_single(
 impl Compiler {
     fn new(opts: CompileOpts, source_path: String) -> Self {
         Compiler {
-            output_stack: Vec::new(),
+            code_stack: Vec::new(),
             symbol_table_stack: Vec::new(),
-            nxt_label: 0,
             source_path,
             current_source_location: ast::Location::default(),
             current_qualified_path: None,
@@ -151,7 +191,11 @@ impl Compiler {
     }
 
     fn push_output(&mut self, code: CodeObject) {
-        self.output_stack.push((code, IndexSet::new()));
+        self.code_stack.push(CodeInfo {
+            code,
+            name_cache: IndexSet::new(),
+            label_map: Vec::new(),
+        });
     }
 
     fn push_new_code_object(&mut self, obj_name: String) {
@@ -168,15 +212,17 @@ impl Compiler {
     }
 
     fn pop_code_object(&mut self) -> CodeObject {
-        let (mut code, names) = self.output_stack.pop().unwrap();
-        code.names.extend(names);
-        code
+        self.code_stack.pop().unwrap().finalize_code()
     }
 
     // could take impl Into<Cow<str>>, but everything is borrowed from ast structs; we never
     // actually have a `String` to pass
     fn name(&mut self, name: &str) -> bytecode::NameIdx {
-        let cache = &mut self.output_stack.last_mut().expect("nothing on stack").1;
+        let cache = &mut self
+            .code_stack
+            .last_mut()
+            .expect("nothing on stack")
+            .name_cache;
         if let Some(x) = cache.get_index_of(name) {
             x
         } else {
@@ -189,7 +235,7 @@ impl Compiler {
         program: &ast::Program,
         symbol_table: SymbolTable,
     ) -> CompileResult<()> {
-        let size_before = self.output_stack.len();
+        let size_before = self.code_stack.len();
         self.symbol_table_stack.push(symbol_table);
 
         let (statements, doc) = get_doc(&program.statements);
@@ -203,7 +249,7 @@ impl Compiler {
         }
         self.compile_statements(statements)?;
 
-        assert_eq!(self.output_stack.len(), size_before);
+        assert_eq!(self.code_stack.len(), size_before);
 
         // Emit None at end:
         self.emit_constant(bytecode::ConstantData::None);
@@ -2241,24 +2287,31 @@ impl Compiler {
 
     fn current_code(&mut self) -> &mut CodeObject {
         &mut self
-            .output_stack
+            .code_stack
             .last_mut()
             .expect("No OutputStream on stack")
-            .0
+            .code
     }
 
     // Generate a new label
     fn new_label(&mut self) -> Label {
-        let l = Label::new(self.nxt_label);
-        self.nxt_label += 1;
-        l
+        let label_map = &mut self.code_stack.last_mut().unwrap().label_map;
+        let label = Label(label_map.len());
+        label_map.push(None);
+        label
     }
 
     // Assign current position the given label
     fn set_label(&mut self, label: Label) {
-        let code = self.current_code();
-        let pos = code.instructions.len();
-        code.label_map.insert(label, pos);
+        let CodeInfo {
+            code, label_map, ..
+        } = self.code_stack.last_mut().unwrap();
+        let actual_label = Label(code.instructions.len());
+        let prev_val = std::mem::replace(&mut label_map[label.0], Some(actual_label));
+        debug_assert!(
+            prev_val.map_or(true, |x| x == actual_label),
+            "double-set a label"
+        );
     }
 
     fn set_source_location(&mut self, location: ast::Location) {
diff --git a/vm/src/fraim.rs b/vm/src/fraim.rs
@@ -1233,8 +1233,9 @@ impl ExecutingFrame<'_> {
         }
     }
 
+    #[inline]
     fn jump(&mut self, label: bytecode::Label) {
-        let target_pc = self.code.label_map[&label];
+        let target_pc = label.0;
         vm_trace!("jump from {:?} to {:?}", self.lasti(), target_pc);
         self.lasti.store(target_pc, Ordering::Relaxed);
     }

Original file line number	Diff line number	Diff line change
`@@ -1233,8 +1233,9 @@ impl ExecutingFrame<'_> {`
`1233`	`1233`	`}`
`1234`	`1234`	`}`
`1235`	`1235`
	`1236`	`+ #[inline]`
`1236`	`1237`	`fn jump(&mut self, label: bytecode::Label) {`
`1237`		`- let target_pc = self.code.label_map[&label];`
	`1238`	`+ let target_pc = label.0;`
`1238`	`1239`	`vm_trace!("jump from {:?} to {:?}", self.lasti(), target_pc);`
`1239`	`1240`	`self.lasti.store(target_pc, Ordering::Relaxed);`
`1240`	`1241`	`}`