star_unpack_helper · RustPython/RustPython@5efbe66
@@ -5,6 +5,8 @@
55//! <https://github.com/python/cpython/blob/main/Python/compile.c>
66//! <https://github.com/micropython/micropython/blob/master/py/compile.c>
778+// spell-checker:ignore starunpack subscripter
9+810#![deny(clippy::cast_possible_truncation)]
9111012use crate::{
@@ -47,9 +49,9 @@ use num_complex::Complex;
4749use num_traits::{Num, ToPrimitive};
4850use ruff_python_ast::{
4951Alias, Arguments, BoolOp, CmpOp, Comprehension, ConversionFlag, DebugText, Decorator, DictItem,
50-ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprFString,
51-ExprList, ExprName, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp, FString,
52-FStringElement, FStringElements, FStringFlags, FStringPart, Identifier, Int, Keyword,
52+ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprContext,
53+ExprFString, ExprList, ExprName, ExprSlice, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp,
54+FString, FStringElement, FStringElements, FStringFlags, FStringPart, Identifier, Int, Keyword,
5355MatchCase, ModExpression, ModModule, Operator, Parameters, Pattern, PatternMatchAs,
5456PatternMatchClass, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton,
5557PatternMatchStar, PatternMatchValue, Singleton, Stmt, StmtExpr, TypeParam, TypeParamParamSpec,
@@ -372,7 +374,158 @@ impl<'src> Compiler<'src> {
372374}
373375}
374376377+/// Type of collection to build in starunpack_helper
378+#[derive(Debug, Clone, Copy, PartialEq)]
379+enum CollectionType {
380+Tuple,
381+List,
382+Set,
383+}
384+375385impl Compiler<'_> {
386+/// Check if the slice is a two-element slice (no step)
387+ // = is_two_element_slice
388+fn is_two_element_slice(slice: &Expr) -> bool {
389+matches!(slice, Expr::Slice(s) if s.step.is_none())
390+}
391+392+/// Compile a slice expression
393+ // = compiler_slice
394+fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult<u32> {
395+// Compile lower
396+if let Some(lower) = &s.lower {
397+self.compile_expression(lower)?;
398+} else {
399+self.emit_load_const(ConstantData::None);
400+}
401+402+// Compile upper
403+if let Some(upper) = &s.upper {
404+self.compile_expression(upper)?;
405+} else {
406+self.emit_load_const(ConstantData::None);
407+}
408+409+// Compile step if present
410+if let Some(step) = &s.step {
411+self.compile_expression(step)?;
412+Ok(3) // Three values on stack
413+} else {
414+Ok(2) // Two values on stack
415+}
416+}
417+418+/// Compile a subscript expression
419+ // = compiler_subscript
420+fn compile_subscript(
421+&mut self,
422+value: &Expr,
423+slice: &Expr,
424+ctx: ExprContext,
425+) -> CompileResult<()> {
426+// 1. Check subscripter and index for Load context
427+// 2. VISIT value
428+// 3. Handle two-element slice specially
429+// 4. Otherwise VISIT slice and emit appropriate instruction
430+431+// For Load context, CPython does some checks (we skip for now)
432+// if ctx == ExprContext::Load {
433+// check_subscripter(value);
434+// check_index(value, slice);
435+// }
436+437+// VISIT(c, expr, e->v.Subscript.value)
438+self.compile_expression(value)?;
439+440+// Handle two-element slice (for Load/Store, not Del)
441+if Self::is_two_element_slice(slice) && !matches!(ctx, ExprContext::Del) {
442+let n = match slice {
443+Expr::Slice(s) => self.compile_slice(s)?,
444+ _ => unreachable!("is_two_element_slice should only return true for Expr::Slice"),
445+};
446+match ctx {
447+ExprContext::Load => {
448+// CPython uses BINARY_SLICE
449+emit!(self, Instruction::BuildSlice { step: n == 3 });
450+emit!(self, Instruction::Subscript);
451+}
452+ExprContext::Store => {
453+// CPython uses STORE_SLICE
454+emit!(self, Instruction::BuildSlice { step: n == 3 });
455+emit!(self, Instruction::StoreSubscript);
456+}
457+ _ => unreachable!(),
458+}
459+} else {
460+// VISIT(c, expr, e->v.Subscript.slice)
461+self.compile_expression(slice)?;
462+463+// Emit appropriate instruction based on context
464+match ctx {
465+ExprContext::Load => emit!(self, Instruction::Subscript),
466+ExprContext::Store => emit!(self, Instruction::StoreSubscript),
467+ExprContext::Del => emit!(self, Instruction::DeleteSubscript),
468+ExprContext::Invalid => {
469+return Err(self.error(CodegenErrorType::SyntaxError(
470+"Invalid expression context".to_owned(),
471+)));
472+}
473+}
474+}
475+476+Ok(())
477+}
478+479+/// Helper function for compiling tuples/lists/sets with starred expressions
480+ ///
481+ /// Parameters:
482+ /// - elts: The elements to compile
483+ /// - pushed: Number of items already on the stack
484+ /// - collection_type: What type of collection to build (tuple, list, set)
485+ ///
486+ // = starunpack_helper in compile.c
487+fn starunpack_helper(
488+&mut self,
489+elts: &[Expr],
490+pushed: u32,
491+collection_type: CollectionType,
492+) -> CompileResult<()> {
493+// Use RustPython's existing approach with BuildXFromTuples
494+let (size, unpack) = self.gather_elements(pushed, elts)?;
495+496+if unpack {
497+// Has starred elements
498+match collection_type {
499+CollectionType::Tuple => {
500+if size > 1 || pushed > 0 {
501+emit!(self, Instruction::BuildTupleFromTuples { size });
502+}
503+// If size == 1 and pushed == 0, the single tuple is already on the stack
504+}
505+CollectionType::List => {
506+emit!(self, Instruction::BuildListFromTuples { size });
507+}
508+CollectionType::Set => {
509+emit!(self, Instruction::BuildSetFromTuples { size });
510+}
511+}
512+} else {
513+// No starred elements
514+match collection_type {
515+CollectionType::Tuple => {
516+emit!(self, Instruction::BuildTuple { size });
517+}
518+CollectionType::List => {
519+emit!(self, Instruction::BuildList { size });
520+}
521+CollectionType::Set => {
522+emit!(self, Instruction::BuildSet { size });
523+}
524+}
525+}
526+527+Ok(())
528+}
376529fn error(&mut self, error: CodegenErrorType) -> CodegenError {
377530self.error_ranged(error, self.current_source_range)
378531}
@@ -1578,10 +1731,10 @@ impl Compiler<'_> {
15781731let idx = self.name(attr.as_str());
15791732emit!(self, Instruction::DeleteAttr { idx });
15801733}
1581-Expr::Subscript(ExprSubscript { value, slice, .. }) => {
1582-self.compile_expression(value)?;
1583- self.compile_expression(slice)?;
1584-emit!(self, Instruction::DeleteSubscript);
1734+Expr::Subscript(ExprSubscript {
1735+ value, slice, ctx, ..
1736+}) => {
1737+self.compile_subscript(value, slice, *ctx)?;
15851738}
15861739Expr::Tuple(ExprTuple { elts, .. }) | Expr::List(ExprList { elts, .. }) => {
15871740for element in elts {
@@ -3946,10 +4099,10 @@ impl Compiler<'_> {
39464099fn compile_store(&mut self, target: &Expr) -> CompileResult<()> {
39474100match &target {
39484101Expr::Name(ExprName { id, .. }) => self.store_name(id.as_str())?,
3949-Expr::Subscript(ExprSubscript { value, slice, .. }) => {
3950-self.compile_expression(value)?;
3951- self.compile_expression(slice)?;
3952-emit!(self, Instruction::StoreSubscript);
4102+Expr::Subscript(ExprSubscript {
4103+ value, slice, ctx, ..
4104+}) => {
4105+self.compile_subscript(value, slice, *ctx)?;
39534106}
39544107Expr::Attribute(ExprAttribute { value, attr, .. }) => {
39554108self.check_forbidden_name(attr.as_str(), NameUsage::Store)?;
@@ -4030,7 +4183,14 @@ impl Compiler<'_> {
40304183self.compile_name(id, NameUsage::Load)?;
40314184AugAssignKind::Name { id }
40324185}
4033-Expr::Subscript(ExprSubscript { value, slice, .. }) => {
4186+Expr::Subscript(ExprSubscript {
4187+ value,
4188+ slice,
4189+ctx: _,
4190+ ..
4191+}) => {
4192+// For augmented assignment, we need to load the value first
4193+// But we can't use compile_subscript directly because we need DUP_TOP2
40344194self.compile_expression(value)?;
40354195self.compile_expression(slice)?;
40364196emit!(self, Instruction::Duplicate2);
@@ -4264,10 +4424,10 @@ impl Compiler<'_> {
42644424// Perform operation:
42654425self.compile_op(op, false);
42664426}
4267-Expr::Subscript(ExprSubscript { value, slice, .. }) => {
4268-self.compile_expression(value)?;
4269- self.compile_expression(slice)?;
4270-emit!(self, Instruction::Subscript);
4427+Expr::Subscript(ExprSubscript {
4428+ value, slice, ctx, ..
4429+}) => {
4430+self.compile_subscript(value, slice, *ctx)?;
42714431}
42724432Expr::UnaryOp(ExprUnaryOp { op, operand, .. }) => {
42734433self.compile_expression(operand)?;
@@ -4298,30 +4458,13 @@ impl Compiler<'_> {
42984458// self.emit_load_const(compile_constant(value));
42994459// }
43004460Expr::List(ExprList { elts, .. }) => {
4301-let (size, unpack) = self.gather_elements(0, elts)?;
4302-if unpack {
4303-emit!(self, Instruction::BuildListFromTuples { size });
4304-} else {
4305-emit!(self, Instruction::BuildList { size });
4306-}
4461+self.starunpack_helper(elts, 0, CollectionType::List)?;
43074462}
43084463Expr::Tuple(ExprTuple { elts, .. }) => {
4309-let (size, unpack) = self.gather_elements(0, elts)?;
4310-if unpack {
4311-if size > 1 {
4312-emit!(self, Instruction::BuildTupleFromTuples { size });
4313-}
4314-} else {
4315-emit!(self, Instruction::BuildTuple { size });
4316-}
4464+self.starunpack_helper(elts, 0, CollectionType::Tuple)?;
43174465}
43184466Expr::Set(ExprSet { elts, .. }) => {
4319-let (size, unpack) = self.gather_elements(0, elts)?;
4320-if unpack {
4321-emit!(self, Instruction::BuildSetFromTuples { size });
4322-} else {
4323-emit!(self, Instruction::BuildSet { size });
4324-}
4467+self.starunpack_helper(elts, 0, CollectionType::Set)?;
43254468}
43264469Expr::Dict(ExprDict { items, .. }) => {
43274470self.compile_dict(items)?;