Fix label issue

src/lib.rs

@@ -8,6 +8,6 @@ extern crate log;
 
 pub mod compile;
 pub mod error;
-pub mod symboltable;
 pub mod output_stream;
 pub mod peephole;
+pub mod symboltable;

src/peephole.rs

@@ -1,12 +1,31 @@
+use crate::compile::Label;
 use crate::output_stream::OutputStream;
 use arrayvec::ArrayVec;
 use rustpython_bytecode::bytecode::{self, CodeObject, Instruction, Location};
 
-const PEEPHOLE_BUFFER_SIZE: usize = 10;
+const PEEPHOLE_BUFFER_SIZE: usize = 20;
+
+struct InstructionMetadata {
+    loc: Location,
+    labels: Vec<Label>,
+}
+
+impl InstructionMetadata {
+    fn from_multiple(metas: Vec<Self>) -> Self {
+        debug_assert!(!metas.is_empty(), "`metas` must not be empty");
+        InstructionMetadata {
+            loc: metas[0].loc.clone(),
+            labels: metas
+                .into_iter()
+                .flat_map(|meta| meta.labels.into_iter())
+                .collect(),
+        }
+    }
+}
 
 pub struct PeepholeOptimizer<O: OutputStream> {
     inner: O,
-    buffer: ArrayVec<[(Instruction, Location); PEEPHOLE_BUFFER_SIZE]>,
+    buffer: ArrayVec<[(Instruction, InstructionMetadata); PEEPHOLE_BUFFER_SIZE]>,
 }
 
 impl<O: OutputStream> From<CodeObject> for PeepholeOptimizer<O> {
@@ -29,25 +48,31 @@ impl<O: OutputStream> PeepholeOptimizer<O> {
         }
     }
 
-    fn emit(&mut self, instruction: Instruction, loc: Location) {
+    fn inner_emit(inner: &mut O, instruction: Instruction, meta: InstructionMetadata) {
-        if self.buffer.is_full() {
+        inner.emit(instruction, meta.loc);
-            let (instr, loc) = self.buffer.remove(0);
+        for label in meta.labels {
-            self.inner.emit(instr, loc);
+            inner.set_label(label);
-            assert_eq!(self.buffer.len(), PEEPHOLE_BUFFER_SIZE - 1)
         }
-        // safe because we just checked that: if full then remove one element from it
-        unsafe { self.buffer.push_unchecked((instruction, loc)) };
     }
 
-    fn pop(&mut self) -> (Instruction, Location) {
+    fn emit(&mut self, instruction: Instruction, meta: InstructionMetadata) {
+        if self.buffer.is_full() {
+            let (instr, meta) = self.buffer.remove(0);
+            Self::inner_emit(&mut self.inner, instr, meta);
+        }
+        // safe because we just checked that: if full then remove one element from it
+        unsafe { self.buffer.push_unchecked((instruction, meta)) };
+    }
+
+    fn pop(&mut self) -> (Instruction, InstructionMetadata) {
         self.buffer.pop().unwrap()
     }
 
-    fn optimize(&mut self, instruction: Instruction, loc: Location) {
+    fn optimize(&mut self, instruction: Instruction, meta: InstructionMetadata) {
         match instruction {
             Instruction::BinaryOperation { op, inplace } => {
-                let (rhs, rhs_loc) = self.pop();
+                let (rhs, rhs_meta) = self.pop();
-                let (lhs, lhs_loc) = self.pop();
+                let (lhs, lhs_meta) = self.pop();
                 macro_rules! lc {
                     ($name:ident {$($field:tt)*}) => {
                         Instruction::LoadConst {
@@ -59,8 +84,11 @@ impl<O: OutputStream> PeepholeOptimizer<O> {
                     };
                 }
                 macro_rules! emitconst {
-                    ($($arg:tt)*) => {
+                    ([$($metas:expr),*], $($arg:tt)*) => {
-                        self.emit(lc!($($arg)*), lhs_loc)
+                        self.emit(
+                            lc!($($arg)*),
+                            InstructionMetadata::from_multiple(vec![$($metas),*]),
+                        )
                     };
                 }
                 macro_rules! op {
@@ -70,40 +98,44 @@ impl<O: OutputStream> PeepholeOptimizer<O> {
                 }
                 match (op, lhs, rhs) {
                     (op!(Add), lc!(Integer, lhs), lc!(Integer, rhs)) => {
-                        emitconst!(Integer, lhs + rhs)
+                        emitconst!([lhs_meta, rhs_meta], Integer, lhs + rhs)
                     }
                     (op!(Subtract), lc!(Integer, lhs), lc!(Integer, rhs)) => {
-                        emitconst!(Integer, lhs - rhs)
+                        emitconst!([lhs_meta, rhs_meta], Integer, lhs - rhs)
+                    }
+                    (op!(Add), lc!(Float, lhs), lc!(Float, rhs)) => {
+                        emitconst!([lhs_meta, rhs_meta], Float, lhs + rhs)
                     }
-                    (op!(Add), lc!(Float, lhs), lc!(Float, rhs)) => emitconst!(Float, lhs + rhs),
                     (op!(Subtract), lc!(Float, lhs), lc!(Float, rhs)) => {
-                        emitconst!(Float, lhs - rhs)
+                        emitconst!([lhs_meta, rhs_meta], Float, lhs - rhs)
                     }
                     (op!(Multiply), lc!(Float, lhs), lc!(Float, rhs)) => {
-                        emitconst!(Float, lhs * rhs)
+                        emitconst!([lhs_meta, rhs_meta], Float, lhs * rhs)
+                    }
+                    (op!(Divide), lc!(Float, lhs), lc!(Float, rhs)) => {
+                        emitconst!([lhs_meta, rhs_meta], Float, lhs / rhs)
                     }
-                    (op!(Divide), lc!(Float, lhs), lc!(Float, rhs)) => emitconst!(Float, lhs / rhs),
                     (op!(Power), lc!(Float, lhs), lc!(Float, rhs)) => {
-                        emitconst!(Float, lhs.powf(rhs))
+                        emitconst!([lhs_meta, rhs_meta], Float, lhs.powf(rhs))
                     }
                     (op!(Add), lc!(String, mut lhs), lc!(String, rhs)) => {
                         lhs.push_str(&rhs);
-                        emitconst!(String, lhs);
+                        emitconst!([lhs_meta, rhs_meta], String, lhs);
                     }
                     (op, lhs, rhs) => {
-                        self.emit(lhs, lhs_loc);
+                        self.emit(lhs, lhs_meta);
-                        self.emit(rhs, rhs_loc);
+                        self.emit(rhs, rhs_meta);
-                        self.emit(Instruction::BinaryOperation { op, inplace }, loc);
+                        self.emit(Instruction::BinaryOperation { op, inplace }, meta);
                     }
                 }
             }
-            other => self.emit(other, loc),
+            other => self.emit(other, meta),
         }
     }
 
     fn flush(&mut self) {
-        for (instruction, location) in self.buffer.drain(..) {
+        for (instruction, meta) in self.buffer.drain(..) {
-            self.inner.emit(instruction, location);
+            Self::inner_emit(&mut self.inner, instruction, meta);
         }
     }
 }
@@ -112,12 +144,19 @@ impl<O> OutputStream for PeepholeOptimizer<O>
 where
     O: OutputStream,
 {
-    fn emit(&mut self, instruction: Instruction, location: Location) {
+    fn emit(&mut self, instruction: Instruction, loc: Location) {
-        self.optimize(instruction, location);
+        self.optimize(
+            instruction,
+            InstructionMetadata {
+                loc,
+                labels: Vec::new(),
+            },
+        );
     }
     fn set_label(&mut self, label: crate::compile::Label) {
-        self.flush();
+        if let Some(instr) = self.buffer.last_mut() {
-        self.inner.set_label(label);
+            instr.1.labels.push(label)
+        }
     }
     fn mark_generator(&mut self) {
         self.inner.mark_generator()