RustPython · youknowone · Aug 10, 2022
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -11,6 +11,7 @@
             "cargo": {
                 "args": [
                     "build",
+                    "--no-default-features",
                     "--package=rustpython"
                 ],
             },

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bytecode/src/lib.rs b/bytecode/src/lib.rs
@@ -246,12 +246,6 @@ pub enum Instruction {
     Duplicate,
     Duplicate2,
     GetIter,
-    Continue {
-        target: Label,
-    },
-    Break {
-        target: Label,
-    },
     Jump {
         target: Label,
     },
@@ -302,9 +296,6 @@ pub enum Instruction {
     YieldValue,
     YieldFrom,
     SetupAnnotation,
-    SetupLoop {
-        break_target: Label,
-    },
 
     /// Setup a finally handler, which will be called whenever one of this events occurs:
     /// - the block is popped
@@ -892,9 +883,7 @@ impl Instruction {
             | SetupFinally { handler: l }
             | SetupExcept { handler: l }
             | SetupWith { end: l }
-            | SetupAsyncWith { end: l }
-            | SetupLoop { break_target: l }
-            | Continue { target: l } => Some(l),
+            | SetupAsyncWith { end: l } => Some(l),
             _ => None,
         }
     }
@@ -912,9 +901,22 @@ impl Instruction {
             | SetupFinally { handler: l }
             | SetupExcept { handler: l }
             | SetupWith { end: l }
-            | SetupAsyncWith { end: l }
-            | SetupLoop { break_target: l }
-            | Continue { target: l } => Some(l),
+            | SetupAsyncWith { end: l } => Some(l),
+            _ => None,
+        }
+    }
+
+    #[inline]
+    pub fn jump_target(&self) -> Option<&Label> {
+        match self {
+            Jump { target: l }
+            | JumpIfTrue { target: l }
+            | JumpIfFalse { target: l }
+            | ForIter { target: l }
+            | SetupFinally { handler: l }
+            | SetupExcept { handler: l }
+            | SetupWith { end: l }
+            | SetupAsyncWith { end: l } => Some(l),
             _ => None,
         }
     }
@@ -930,10 +932,7 @@ impl Instruction {
     /// assert!(jump_inst.unconditional_branch())
     /// ```
     pub fn unconditional_branch(&self) -> bool {
-        matches!(
-            self,
-            Jump { .. } | Continue { .. } | Break { .. } | ReturnValue | Raise { .. }
-        )
+        matches!(self, Jump { .. } | ReturnValue | Raise { .. })
     }
 
     /// What effect this instruction has on the stack
@@ -974,8 +973,6 @@ impl Instruction {
             Duplicate => 1,
             Duplicate2 => 2,
             GetIter => 0,
-            Continue { .. } => 0,
-            Break { .. } => 0,
             Jump { .. } => 0,
             JumpIfTrue { .. } | JumpIfFalse { .. } => -1,
             JumpIfTrueOrPop { .. } | JumpIfFalseOrPop { .. } => {
@@ -1009,11 +1006,7 @@ impl Instruction {
             ReturnValue => -1,
             YieldValue => 0,
             YieldFrom => -1,
-            SetupAnnotation
-            | SetupLoop { .. }
-            | SetupFinally { .. }
-            | EnterFinally
-            | EndFinally => 0,
+            SetupAnnotation | SetupFinally { .. } | EnterFinally | EndFinally => 0,
             SetupExcept { .. } => {
                 if jump {
                     1
@@ -1161,8 +1154,6 @@ impl Instruction {
             Duplicate => w!(Duplicate),
             Duplicate2 => w!(Duplicate2),
             GetIter => w!(GetIter),
-            Continue { target } => w!(Continue, target),
-            Break { target } => w!(Break, target),
             Jump { target } => w!(Jump, target),
             JumpIfTrue { target } => w!(JumpIfTrue, target),
             JumpIfFalse { target } => w!(JumpIfFalse, target),
@@ -1181,7 +1172,6 @@ impl Instruction {
             YieldValue => w!(YieldValue),
             YieldFrom => w!(YieldFrom),
             SetupAnnotation => w!(SetupAnnotation),
-            SetupLoop { break_target } => w!(SetupLoop, break_target),
             SetupExcept { handler } => w!(SetupExcept, handler),
             SetupFinally { handler } => w!(SetupFinally, handler),
             EnterFinally => w!(EnterFinally),

diff --git a/compiler/Cargo.toml b/compiler/Cargo.toml
@@ -16,6 +16,7 @@ num-complex = { version = "0.4.0", features = ["serde"] }
 num-traits = "0.2.14"
 log = "0.4.16"
 ahash = "0.7.6"
+optional = "0.5.0"
 
 [dev-dependencies]
 rustpython-parser = { path = "../parser" }

diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs
@@ -763,15 +763,16 @@ impl Compiler {
             }
             Break => match self.ctx.loop_data {
                 Some((_, end)) => {
-                    self.emit(Instruction::Break { target: end });
+                    self.emit(Instruction::Pop);
+                    self.emit(Instruction::Jump { target: end });
                 }
                 None => {
                     return Err(self.error_loc(CompileErrorType::InvalidBreak, statement.location));
                 }
             },
             Continue => match self.ctx.loop_data {
                 Some((start, _)) => {
-                    self.emit(Instruction::Continue { target: start });
+                    self.emit(Instruction::Jump { target: start });
                 }
                 None => {
                     return Err(
@@ -1367,9 +1368,6 @@ impl Compiler {
         let else_block = self.new_block();
         let after_block = self.new_block();
 
-        self.emit(Instruction::SetupLoop {
-            break_target: after_block,
-        });
         self.switch_to_block(while_block);
 
         self.compile_jump_if(test, false, else_block)?;
@@ -1381,7 +1379,6 @@ impl Compiler {
             target: while_block,
         });
         self.switch_to_block(else_block);
-        self.emit(Instruction::PopBlock);
         self.compile_statements(orelse)?;
         self.switch_to_block(after_block);
         Ok(())
@@ -1472,10 +1469,6 @@ impl Compiler {
         let else_block = self.new_block();
         let after_block = self.new_block();
 
-        self.emit(Instruction::SetupLoop {
-            break_target: after_block,
-        });
-
         // The thing iterated:
         self.compile_expression(iter)?;
 
@@ -1511,7 +1504,6 @@ impl Compiler {
         if is_async {
             self.emit(Instruction::EndAsyncFor);
         }
-        self.emit(Instruction::PopBlock);
         self.compile_statements(orelse)?;
 
         self.switch_to_block(after_block);

diff --git a/compiler/src/ir.rs b/compiler/src/ir.rs
@@ -1,4 +1,5 @@
 use crate::IndexSet;
+use optional::Optioned;
 use rustpython_bytecode::{CodeFlags, CodeObject, ConstantData, Instruction, Label, Location};
 
 pub type BlockIdx = Label;
@@ -16,6 +17,13 @@ pub struct InstructionInfo {
 pub struct Block {
     pub instructions: Vec<InstructionInfo>,
     pub next: BlockIdx,
+    // has_return: bool,
+    // precedessors
+    // is_nofallthrough: bool,
+    // has_exit: bool,
+    // is_visited: bool,
+    // startdepth: i32,
+    // offset: usize,
 }
 impl Default for Block {
     fn default() -> Self {
@@ -162,72 +170,75 @@ impl CodeInfo {
     }
 
     fn max_stackdepth(&self) -> u32 {
-        let mut maxdepth = 0u32;
+        let mut maxdepth = 0i32;
         let mut stack = Vec::with_capacity(self.blocks.len());
-        let mut startdepths = vec![u32::MAX; self.blocks.len()];
-        startdepths[0] = 0;
-        stack.push(Label(0));
+        let mut startdepths = vec![Optioned::<i32>::none(); self.blocks.len()];
+        stackdepth_push(&mut stack, &mut startdepths, Label(0), 1);
         const DEBUG: bool = false;
         'process_blocks: while let Some(block) = stack.pop() {
-            let mut depth = startdepths[block.0 as usize];
+            let block_idx = block.0 as usize;
+            let mut depth = startdepths[block_idx].expect("must be set");
             if DEBUG {
                 eprintln!("===BLOCK {}===", block.0);
             }
-            let block = &self.blocks[block.0 as usize];
+            let block = &self.blocks[block_idx];
             for i in &block.instructions {
                 let instr = &i.instr;
+                println!("instr: {instr:?}");
                 let effect = instr.stack_effect(false);
                 if DEBUG {
                     eprint!("{instr:?}: {depth} {effect:+} => ");
                 }
-                let new_depth = add_ui(depth, effect);
+                let new_depth = depth + effect;
                 if DEBUG {
                     eprintln!("{new_depth}");
                 }
                 if new_depth > maxdepth {
                     maxdepth = new_depth
                 }
-                // we don't want to worry about Continue, it uses unwinding to jump to
-                // its targets and as such the stack size is taken care of in frame.rs by setting
-                // it back to the level it was at when SetupLoop was run
-                let jump_label = instr
-                    .label_arg()
-                    .filter(|_| !matches!(instr, Instruction::Continue { .. }));
+                assert!(depth >= 0);
+                let jump_label = instr.jump_target();
                 if let Some(&target_block) = jump_label {
                     let effect = instr.stack_effect(true);
-                    let target_depth = add_ui(depth, effect);
+                    let target_depth = depth + effect;
                     if target_depth > maxdepth {
                         maxdepth = target_depth
                     }
+                    assert!(target_depth >= 0);
                     stackdepth_push(&mut stack, &mut startdepths, target_block, target_depth);
                 }
                 depth = new_depth;
-                if instr.unconditional_branch() {
+                let unconditional = instr.unconditional_branch();
+                if unconditional {
+                    // next = null / break
                     continue 'process_blocks;
                 }
             }
+            // assert!(block.nofallthrough);
             stackdepth_push(&mut stack, &mut startdepths, block.next, depth);
         }
         if DEBUG {
             eprintln!("DONE: {maxdepth}");
         }
-        maxdepth
+        maxdepth as u32
     }
 }
 
-fn stackdepth_push(stack: &mut Vec<Label>, startdepths: &mut [u32], target: Label, depth: u32) {
+fn stackdepth_push(
+    stack: &mut Vec<Label>,
+    startdepths: &mut [Optioned<i32>],
+    target: Label,
+    depth: i32,
+) {
     let block_depth = &mut startdepths[target.0 as usize];
-    if *block_depth == u32::MAX || depth > *block_depth {
-        *block_depth = depth;
-        stack.push(target);
+    if !block_depth.map_or(true, |d| d == depth) {
+        println!("here!");
     }
-}
-
-fn add_ui(a: u32, b: i32) -> u32 {
-    if b < 0 {
-        a - b.wrapping_abs() as u32
-    } else {
-        a + b as u32
+    assert!(block_depth.map_or(true, |d| d == depth));
+    if block_depth.map_or(true, |d| d < std::cmp::min(depth, 100)) {
+        assert!(block_depth.is_none());
+        *block_depth = Optioned::some(depth);
+        stack.push(target);
     }
 }
 

diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
@@ -379,7 +379,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     _ => Err(JitCompileError::NotSupported),
                 }
             }
-            Instruction::SetupLoop { .. } | Instruction::PopBlock => {
+            Instruction::PopBlock => {
                 // TODO: block support
                 Ok(())
             }