Split Rotate into Rotate2 & Rotate3

Author: coolreader18

bytecode/src/lib.rs

@@ -265,9 +265,8 @@ pub enum Instruction {
         op: ComparisonOperator,
     },
     Pop,
-    Rotate {
-        amount: u32,
-    },
+    Rotate2,
+    Rotate3,
     Duplicate,
     Duplicate2,
     GetIter,
@@ -972,7 +971,7 @@ impl Instruction {
             UnaryOperation { .. } => 0,
             BinaryOperation { .. } | BinaryOperationInplace { .. } | CompareOperation { .. } => -1,
             Pop => -1,
-            Rotate { .. } => 0,
+            Rotate2 | Rotate3 => 0,
             Duplicate => 1,
             Duplicate2 => 2,
             GetIter => 0,
@@ -1157,7 +1156,8 @@ impl Instruction {
             LoadAttr { idx } => w!(LoadAttr, name(*idx)),
             CompareOperation { op } => w!(CompareOperation, format_args!("{:?}", op)),
             Pop => w!(Pop),
-            Rotate { amount } => w!(Rotate, amount),
+            Rotate2 => w!(Rotate2),
+            Rotate3 => w!(Rotate3),
             Duplicate => w!(Duplicate),
             Duplicate2 => w!(Duplicate2),
             GetIter => w!(GetIter),

compiler/src/compile.rs

@@ -1122,7 +1122,7 @@ impl Compiler {
 
         self.emit(Instruction::Duplicate);
         self.load_docstring(doc_str);
-        self.emit(Instruction::Rotate { amount: 2 });
+        self.emit(Instruction::Rotate2);
         let doc = self.name("__doc__");
         self.emit(Instruction::StoreAttr { idx: doc });
 
@@ -1498,7 +1498,7 @@ impl Compiler {
             self.compile_expression(val)?;
             // store rhs for the next comparison in chain
             self.emit(Instruction::Duplicate);
-            self.emit(Instruction::Rotate { amount: 3 });
+            self.emit(Instruction::Rotate3);
 
             self.emit(Instruction::CompareOperation {
                 op: compile_cmpop(op),
@@ -1525,7 +1525,7 @@ impl Compiler {
 
             // early exit left us with stack: `rhs, comparison_result`. We need to clean up rhs.
             self.switch_to_block(break_block);
-            self.emit(Instruction::Rotate { amount: 2 });
+            self.emit(Instruction::Rotate2);
             self.emit(Instruction::Pop);
 
             self.switch_to_block(after_block);
@@ -1692,12 +1692,12 @@ impl Compiler {
             }
             AugAssignKind::Subscript => {
                 // stack: CONTAINER SLICE RESULT
-                self.emit(Instruction::Rotate { amount: 3 });
+                self.emit(Instruction::Rotate3);
                 self.emit(Instruction::StoreSubscript);
             }
             AugAssignKind::Attr { idx } => {
                 // stack: CONTAINER RESULT
-                self.emit(Instruction::Rotate { amount: 2 });
+                self.emit(Instruction::Rotate2);
                 self.emit(Instruction::StoreAttr { idx });
             }
         }

jit/tests/common.rs

@@ -132,17 +132,13 @@ impl StackMachine {
                 let value = self.stack.last().unwrap().clone();
                 self.stack.push(value);
             }
-            Instruction::Rotate { amount } => {
-                let mut values = Vec::new();
-
-                // Pop all values from stack:
-                values.extend(self.stack.drain(self.stack.len() - amount as usize..));
-
-                // Push top of stack back first:
-                self.stack.push(values.pop().unwrap());
-
-                // Push other value back in order:
-                self.stack.extend(values);
+            Instruction::Rotate2 => {
+                let i = self.stack.len() - 2;
+                self.stack[i..].rotate_right(1);
+            }
+            Instruction::Rotate3 => {
+                let i = self.stack.len() - 3;
+                self.stack[i..].rotate_right(1);
             }
             Instruction::ReturnValue => return true,
             _ => unimplemented!(

vm/src/frame.rs

@@ -623,7 +623,11 @@ impl ExecutingFrame<'_> {
                 self.push_value(top);
                 Ok(None)
             }
-            bytecode::Instruction::Rotate { amount } => self.execute_rotate(*amount),
+            // splitting the instructions like this offloads the cost of "dynamic" dispatch (on the
+            // amount to rotate) to the opcode dispatcher, and generates optimized code for the
+            // concrete cases we actually have
+            bytecode::Instruction::Rotate2 => self.execute_rotate(2),
+            bytecode::Instruction::Rotate3 => self.execute_rotate(3),
             bytecode::Instruction::BuildString { size } => {
                 let s = self
                     .pop_multiple(*size as usize)
@@ -1248,26 +1252,10 @@ impl ExecutingFrame<'_> {
         }
     }
 
-    fn execute_rotate(&mut self, amount: u32) -> FrameResult {
-        // Shuffles top of stack amount down
-        if amount < 2 {
-            self.fatal("Can only rotate two or more values");
-        }
-
-        let mut values = Vec::new();
-
-        // Pop all values from stack:
-        for _ in 0..amount {
-            values.push(self.pop_value());
-        }
-
-        // Push top of stack back first:
-        self.push_value(values.remove(0));
-
-        // Push other value back in order:
-        for value in values.into_iter().rev() {
-            self.push_value(value);
-        }
+    #[inline(always)]
+    fn execute_rotate(&mut self, amount: usize) -> FrameResult {
+        let i = self.state.stack.len() - amount;
+        self.state.stack[i..].rotate_right(1);
         Ok(None)
     }