From 8b0e806e39591887228e7a2464f340030aa973bf Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 23 Nov 2020 20:10:46 -0600
Subject: [PATCH] Fix rustpython-jit to not use label_map

---
 bytecode/src/bytecode.rs | 17 +++++++++++------
 jit/src/instructions.rs  | 21 +++++++++++++--------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
index f13427bee..33c673e35 100644
--- a/bytecode/src/bytecode.rs
+++ b/bytecode/src/bytecode.rs
@@ -595,12 +595,7 @@ impl<C: Constant> CodeObject<C> {
         }
     }
 
-    fn display_inner(
-        &self,
-        f: &mut fmt::Formatter,
-        expand_codeobjects: bool,
-        level: usize,
-    ) -> fmt::Result {
+    pub fn label_targets(&self) -> BTreeSet<Label> {
         let mut label_targets = BTreeSet::new();
         for instruction in &self.instructions {
             let label = match instruction {
@@ -623,6 +618,16 @@ impl<C: Constant> CodeObject<C> {
             };
             label_targets.insert(*label);
         }
+        label_targets
+    }
+
+    fn display_inner(
+        &self,
+        f: &mut fmt::Formatter,
+        expand_codeobjects: bool,
+        level: usize,
+    ) -> fmt::Result {
+        let label_targets = self.label_targets();
 
         for (offset, instruction) in self.instructions.iter().enumerate() {
             let arrow = if label_targets.contains(&Label(offset)) {
diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
index 584970298..2cb89fb4c 100644
--- a/jit/src/instructions.rs
+++ b/jit/src/instructions.rs
@@ -103,11 +103,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         }
     }
 
-    fn get_or_create_block(&mut self, label: &Label) -> Block {
+    fn get_or_create_block(&mut self, label: Label) -> Block {
         let builder = &mut self.builder;
         *self
             .label_to_block
-            .entry(*label)
+            .entry(label)
             .or_insert_with(|| builder.create_block())
     }
 
@@ -115,11 +115,16 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         &mut self,
         bytecode: &CodeObject<C>,
     ) -> Result<(), JitCompileError> {
-        let offset_to_label: HashMap<&usize, &Label> =
-            bytecode.label_map.iter().map(|(k, v)| (v, k)).collect();
+        // TODO: figure out if this is sufficient -- previously individual labels were associated
+        // pretty much per-bytecode that uses them, or at least per "type" of block -- in theory an
+        // if block and a with block might jump to the same place. Now it's all "flattened", so
+        // there might be less distinction between different types of blocks going off
+        // label_targets alone
+        let label_targets = bytecode.label_targets();
 
         for (offset, instruction) in bytecode.instructions.iter().enumerate() {
-            if let Some(&label) = offset_to_label.get(&offset) {
+            let label = Label(offset);
+            if label_targets.contains(&label) {
                 let block = self.get_or_create_block(label);
 
                 // If the current block is not terminated/filled just jump
@@ -190,7 +195,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
 
                 let val = self.boolean_val(cond)?;
-                let then_block = self.get_or_create_block(target);
+                let then_block = self.get_or_create_block(*target);
                 self.builder.ins().brz(val, then_block, &[]);
 
                 let block = self.builder.create_block();
@@ -203,7 +208,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
 
                 let val = self.boolean_val(cond)?;
-                let then_block = self.get_or_create_block(target);
+                let then_block = self.get_or_create_block(*target);
                 self.builder.ins().brnz(val, then_block, &[]);
 
                 let block = self.builder.create_block();
@@ -213,7 +218,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 Ok(())
             }
             Instruction::Jump { target } => {
-                let target_block = self.get_or_create_block(target);
+                let target_block = self.get_or_create_block(*target);
                 self.builder.ins().jump(target_block, &[]);
 
                 Ok(())