diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 4a4f17edcfa..387c2e819ed 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -35,8 +35,8 @@ use ruff_text_size::{Ranged, TextRange}; use rustpython_compiler_core::{ Mode, OneIndexed, PositionEncoding, SourceFile, SourceLocation, bytecode::{ - self, Arg as OpArgMarker, BinaryOperator, CodeObject, ComparisonOperator, ConstantData, - Instruction, Invert, OpArg, OpArgType, UnpackExArgs, + self, Arg as OpArgMarker, BinaryOperator, BuildSliceArgCount, CodeObject, + ComparisonOperator, ConstantData, Instruction, Invert, OpArg, OpArgType, UnpackExArgs, }, }; use rustpython_wtf8::Wtf8Buf; @@ -401,7 +401,7 @@ impl Compiler { /// Compile a slice expression // = compiler_slice - fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult { + fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult { // Compile lower if let Some(lower) = &s.lower { self.compile_expression(lower)?; @@ -416,13 +416,14 @@ impl Compiler { self.emit_load_const(ConstantData::None); } - // Compile step if present - if let Some(step) = &s.step { - self.compile_expression(step)?; - Ok(3) // Three values on stack - } else { - Ok(2) // Two values on stack - } + Ok(match &s.step { + Some(step) => { + // Compile step if present + self.compile_expression(step)?; + BuildSliceArgCount::Three + } + None => BuildSliceArgCount::Two, + }) } /// Compile a subscript expression @@ -449,19 +450,19 @@ impl Compiler { // Handle two-element slice (for Load/Store, not Del) if Self::is_two_element_slice(slice) && !matches!(ctx, ExprContext::Del) { - let n = match slice { + let argc = match slice { Expr::Slice(s) => self.compile_slice(s)?, _ => unreachable!("is_two_element_slice should only return true for Expr::Slice"), }; match ctx { ExprContext::Load => { // CPython uses BINARY_SLICE - emit!(self, Instruction::BuildSlice { step: n == 3 }); + emit!(self, Instruction::BuildSlice { argc }); emit!(self, Instruction::Subscript); } ExprContext::Store => { // CPython uses STORE_SLICE - emit!(self, Instruction::BuildSlice { step: n == 3 }); + emit!(self, Instruction::BuildSlice { argc }); emit!(self, Instruction::StoreSubscript); } _ => unreachable!(), @@ -4587,8 +4588,11 @@ impl Compiler { if let Some(step) = step { self.compile_expression(step)?; } - let step = step.is_some(); - emit!(self, Instruction::BuildSlice { step }); + let argc = match step { + Some(_) => BuildSliceArgCount::Three, + None => BuildSliceArgCount::Two, + }; + emit!(self, Instruction::BuildSlice { argc }); } Expr::Yield(ExprYield { value, .. }) => { if !self.ctx.in_func() { diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs index 144054860e4..8a095de6114 100644 --- a/crates/compiler-core/src/bytecode.rs +++ b/crates/compiler-core/src/bytecode.rs @@ -10,7 +10,7 @@ use itertools::Itertools; use malachite_bigint::BigInt; use num_complex::Complex64; use rustpython_wtf8::{Wtf8, Wtf8Buf}; -use std::{collections::BTreeSet, fmt, hash, marker::PhantomData, mem, ops::Deref}; +use std::{collections::BTreeSet, fmt, hash, marker::PhantomData, mem, num::NonZeroU8, ops::Deref}; #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)] #[repr(i8)] @@ -760,8 +760,7 @@ pub enum Instruction { index: Arg, }, BuildSlice { - /// whether build a slice with a third step argument - step: Arg, + argc: Arg, }, ListAppend { i: Arg, @@ -1151,6 +1150,48 @@ op_arg_enum!( } ); +/// Specifies if a slice is built with either 2 or 3 arguments. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum BuildSliceArgCount { + /// ```py + /// x[5:10] + /// ``` + Two, + /// ```py + /// x[5:10:2] + /// ``` + Three, +} + +impl OpArgType for BuildSliceArgCount { + #[inline(always)] + fn from_op_arg(x: u32) -> Option { + Some(match x { + 2 => Self::Two, + 3 => Self::Three, + _ => return None, + }) + } + + #[inline(always)] + fn to_op_arg(self) -> u32 { + u32::from(self.argc().get()) + } +} + +impl BuildSliceArgCount { + /// Get the numeric value of `Self`. + #[must_use] + pub const fn argc(self) -> NonZeroU8 { + let inner = match self { + Self::Two => 2, + Self::Three => 3, + }; + // Safety: `inner` can be either 2 or 3. + unsafe { NonZeroU8::new_unchecked(inner) } + } +} + #[derive(Copy, Clone)] pub struct UnpackExArgs { pub before: u8, @@ -1547,7 +1588,11 @@ impl Instruction { -(nargs as i32) + 1 } DictUpdate { .. } => -1, - BuildSlice { step } => -2 - (step.get(arg) as i32) + 1, + BuildSlice { argc } => { + // push 1 + // pops either 2/3 + 1 - (argc.get(arg).argc().get() as i32) + } ListAppend { .. } | SetAdd { .. } => -1, MapAdd { .. } => -2, PrintExpr => -1, @@ -1734,7 +1779,7 @@ impl Instruction { BuildMap { size } => w!(BuildMap, size), BuildMapForCall { size } => w!(BuildMapForCall, size), DictUpdate { index } => w!(DictUpdate, index), - BuildSlice { step } => w!(BuildSlice, step), + BuildSlice { argc } => w!(BuildSlice, ?argc), ListAppend { i } => w!(ListAppend, i), SetAdd { i } => w!(SetAdd, i), MapAdd { i } => w!(MapAdd, i), diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 148abacf6ad..088b5dcb4b5 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -834,8 +834,8 @@ impl ExecutingFrame<'_> { dict.merge_object(source, vm)?; Ok(None) } - bytecode::Instruction::BuildSlice { step } => { - self.execute_build_slice(vm, step.get(arg)) + bytecode::Instruction::BuildSlice { argc } => { + self.execute_build_slice(vm, argc.get(arg)) } bytecode::Instruction::ListAppend { i } => { let item = self.pop_value(); @@ -1777,8 +1777,15 @@ impl ExecutingFrame<'_> { Ok(None) } - fn execute_build_slice(&mut self, vm: &VirtualMachine, step: bool) -> FrameResult { - let step = if step { Some(self.pop_value()) } else { None }; + fn execute_build_slice( + &mut self, + vm: &VirtualMachine, + argc: bytecode::BuildSliceArgCount, + ) -> FrameResult { + let step = match argc { + bytecode::BuildSliceArgCount::Two => None, + bytecode::BuildSliceArgCount::Three => Some(self.pop_value()), + }; let stop = self.pop_value(); let start = self.pop_value();