implement SWAP instruction

arihant2math · arihant2math · commit be4df5d86e8c · 2025-01-21T11:31:10.000-08:00
Signed-off-by: Ashwin Naren &lt;arihant2math@gmail.com&gt;
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
@@ -217,31 +217,29 @@ macro_rules! emit {
 }
 
 struct PatternContext {
-    // // A list of strings corresponding to name captures. It is used to track:
-    // // - Repeated name assignments in the same pattern.
-    // // - Different name assignments in alternatives.
-    // // - The order of name assignments in alternatives.
-    // PyObject *stores;
-    // // If 0, any name captures against our subject will raise.
-    // int allow_irrefutable;
-    // // An array of blocks to jump to on failure. Jumping to fail_pop[i] will pop
-    // // i items off of the stack. The end result looks like this (with each block
-    // // falling through to the next):
-    // // fail_pop[4]: POP_TOP
-    // // fail_pop[3]: POP_TOP
-    // // fail_pop[2]: POP_TOP
-    // // fail_pop[1]: POP_TOP
-    // // fail_pop[0]: NOP
-    // jump_target_label *fail_pop;
     // // The current length of fail_pop.
     // Py_ssize_t fail_pop_size;
-    // // The number of items on top of the stack that need to *stay* on top of the
-    // // stack. Variable captures go beneath these. All of them will be popped on
-    // // failure.
-    // Py_ssize_t on_top;
+
+    /// Py_ssize_t on_top;
+    /// A list of strings corresponding to name captures. It is used to track:
+    /// - Repeated name assignments in the same pattern.
+    /// - Different name assignments in alternatives.
+    /// - The order of name assignments in alternatives.
     pub stores: Vec<String>,
+    /// If 0, any name captures against our subject will raise.
     pub allow_irrefutable: bool,
+    /// An array of blocks to jump to on failure. Jumping to fail_pop[i] will pop
+    /// i items off of the stack. The end result looks like this (with each block
+    /// falling through to the next):
+    /// fail_pop[4]: POP_TOP
+    /// fail_pop[3]: POP_TOP
+    /// fail_pop[2]: POP_TOP
+    /// fail_pop[1]: POP_TOP
+    /// fail_pop[0]: NOP
     pub fail_pop: Vec<BlockIdx>,
+    /// The number of items on top of the stack that need to *stay* on top of the
+    /// stack. Variable captures go beneath these. All of them will be popped on
+    /// failure.
     pub on_top: usize,
 }
 
@@ -2055,9 +2053,8 @@ impl Compiler {
         pattern_context: &mut PatternContext,
     ) -> CompileResult<()> {
         self.compile_expression(subject)?;
-        // Blocks at the end of the switch statement that we jump to after finishing a branch
-        // TODO: optimize, we can reuse the same block for all cases
-        let mut end_block = self.new_block();
+        // Block at the end of the switch statement that we jump to after finishing a branch
+        let end_block = self.new_block();
 
         let match_case_type = cases.last().expect("cases is not empty");
         let has_default = match_case_type.pattern.is_match_star() && 1 < cases.len();
@@ -2117,7 +2114,6 @@ impl Compiler {
             self.compile_statements(&m.body)?;
         }
 
-        self.switch_to_block(end_block);
         Ok(())
     }
 
diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
@@ -432,6 +432,7 @@ pub enum Instruction {
         op: Arg<ComparisonOperator>,
     },
     Pop,
+    Swap(Arg<u32>),
     Rotate2,
     Rotate3,
     Duplicate,
@@ -1202,6 +1203,7 @@ impl Instruction {
             | TestOperation { .. }
             | CompareOperation { .. } => -1,
             Pop => -1,
+            Swap(_) => 0,
             Rotate2 | Rotate3 => 0,
             Duplicate => 1,
             Duplicate2 => 2,
@@ -1398,6 +1400,7 @@ impl Instruction {
             TestOperation { op } => w!(TestOperation, ?op),
             CompareOperation { op } => w!(CompareOperation, ?op),
             Pop => w!(Pop),
+            Swap(a) => w!(Swap, a),
             Rotate2 => w!(Rotate2),
             Rotate3 => w!(Rotate3),
             Duplicate => w!(Duplicate),
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -673,6 +673,11 @@ impl ExecutingFrame<'_> {
                 self.pop_value();
                 Ok(None)
             }
+            bytecode::Instruction::Swap(i) => {
+                let len = self.state.stack.len();
+                self.state.stack.swap(len - 1, len - 1 - i.get(arg) as usize);
+                Ok(None)
+            }
             bytecode::Instruction::Duplicate => {
                 // Duplicate top of stack
                 let value = self.top_value();
