Skip to content

Commit 9f54540

Browse files
youknowoneyouknowone
committed
star_unpack_helper
1 parent c0c4852 commit 9f54540

File tree

1 file changed

+182
-35
lines changed

1 file changed

+182
-35
lines changed

‎compiler/codegen/src/compile.rs‎

Lines changed: 182 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -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>
77
8+
// spell-checker:ignore starunpack
9+
810
#![deny(clippy::cast_possible_truncation)]
911

1012
use crate::{
@@ -47,8 +49,8 @@ use num_complex::Complex;
4749
use num_traits::{Num, ToPrimitive};
4850
use ruff_python_ast::{
4951
Alias, Arguments, BoolOp, CmpOp, Comprehension, ConversionFlag, DebugText, Decorator, DictItem,
50-
ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprFString,
51-
ExprList, ExprName, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp, FString,
52+
ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprContext,
53+
ExprFString, ExprList, ExprName, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp, FString,
5254
FStringElement, FStringElements, FStringFlags, FStringPart, Identifier, Int, Keyword,
5355
MatchCase, ModExpression, ModModule, Operator, Parameters, Pattern, PatternMatchAs,
5456
PatternMatchClass, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton,
@@ -372,7 +374,162 @@ impl<'src> Compiler<'src> {
372374
}
373375
}
374376

377+
/// 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+
375385
impl Compiler<'_> {
386+
/// Check if the slice is a two-element slice (no step)
387+
// = is_two_element_slice
388+
fn is_two_element_slice(&self, 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, slice: &Expr) -> CompileResult<u32> {
395+
match slice {
396+
Expr::Slice(s) => {
397+
// Compile lower
398+
if let Some(lower) = &s.lower {
399+
self.compile_expression(lower)?;
400+
} else {
401+
self.emit_load_const(ConstantData::None);
402+
}
403+
404+
// Compile upper
405+
if let Some(upper) = &s.upper {
406+
self.compile_expression(upper)?;
407+
} else {
408+
self.emit_load_const(ConstantData::None);
409+
}
410+
411+
// Compile step if present
412+
if let Some(step) = &s.step {
413+
self.compile_expression(step)?;
414+
Ok(3) // Three values on stack
415+
} else {
416+
Ok(2) // Two values on stack
417+
}
418+
}
419+
_ => Err(self.error(CodegenErrorType::SyntaxError(
420+
"compile_slice expects a slice expression".to_owned(),
421+
))),
422+
}
423+
}
424+
425+
/// Compile a subscript expression
426+
// = compiler_subscript
427+
fn compile_subscript(
428+
&mut self,
429+
value: &Expr,
430+
slice: &Expr,
431+
ctx: ExprContext,
432+
) -> CompileResult<()> {
433+
// 1. Check subscripter and index for Load context
434+
// 2. VISIT value
435+
// 3. Handle two-element slice specially
436+
// 4. Otherwise VISIT slice and emit appropriate instruction
437+
438+
// For Load context, CPython does some checks (we skip for now)
439+
// if ctx == ExprContext::Load {
440+
// check_subscripter(value);
441+
// check_index(value, slice);
442+
// }
443+
444+
// VISIT(c, expr, e->v.Subscript.value)
445+
self.compile_expression(value)?;
446+
447+
// Handle two-element slice (for Load/Store, not Del)
448+
if self.is_two_element_slice(slice) && !matches!(ctx, ExprContext::Del) {
449+
let n = self.compile_slice(slice)?;
450+
match ctx {
451+
ExprContext::Load => {
452+
// CPython uses BINARY_SLICE
453+
emit!(self, Instruction::BuildSlice { step: n == 3 });
454+
emit!(self, Instruction::Subscript);
455+
}
456+
ExprContext::Store => {
457+
// CPython uses STORE_SLICE
458+
emit!(self, Instruction::BuildSlice { step: n == 3 });
459+
emit!(self, Instruction::StoreSubscript);
460+
}
461+
_ => unreachable!(),
462+
}
463+
} else {
464+
// VISIT(c, expr, e->v.Subscript.slice)
465+
self.compile_expression(slice)?;
466+
467+
// Emit appropriate instruction based on context
468+
match ctx {
469+
ExprContext::Load => emit!(self, Instruction::Subscript),
470+
ExprContext::Store => emit!(self, Instruction::StoreSubscript),
471+
ExprContext::Del => emit!(self, Instruction::DeleteSubscript),
472+
ExprContext::Invalid => {
473+
return Err(self.error(CodegenErrorType::SyntaxError(
474+
"Invalid expression context".to_owned(),
475+
)));
476+
}
477+
}
478+
}
479+
480+
Ok(())
481+
}
482+
483+
/// Helper function for compiling tuples/lists/sets with starred expressions
484+
///
485+
/// Parameters:
486+
/// - elts: The elements to compile
487+
/// - pushed: Number of items already on the stack
488+
/// - collection_type: What type of collection to build (tuple, list, set)
489+
///
490+
// = starunpack_helper in compile.c
491+
fn starunpack_helper(
492+
&mut self,
493+
elts: &[Expr],
494+
pushed: u32,
495+
collection_type: CollectionType,
496+
) -> CompileResult<()> {
497+
// Use RustPython's existing approach with BuildXFromTuples
498+
let (size, unpack) = self.gather_elements(pushed, elts)?;
499+
500+
if unpack {
501+
// Has starred elements
502+
match collection_type {
503+
CollectionType::Tuple => {
504+
if size > 1 || pushed > 0 {
505+
emit!(self, Instruction::BuildTupleFromTuples { size });
506+
}
507+
// If size == 1 and pushed == 0, the single tuple is already on the stack
508+
}
509+
CollectionType::List => {
510+
emit!(self, Instruction::BuildListFromTuples { size });
511+
}
512+
CollectionType::Set => {
513+
emit!(self, Instruction::BuildSetFromTuples { size });
514+
}
515+
}
516+
} else {
517+
// No starred elements
518+
match collection_type {
519+
CollectionType::Tuple => {
520+
emit!(self, Instruction::BuildTuple { size });
521+
}
522+
CollectionType::List => {
523+
emit!(self, Instruction::BuildList { size });
524+
}
525+
CollectionType::Set => {
526+
emit!(self, Instruction::BuildSet { size });
527+
}
528+
}
529+
}
530+
531+
Ok(())
532+
}
376533
fn error(&mut self, error: CodegenErrorType) -> CodegenError {
377534
self.error_ranged(error, self.current_source_range)
378535
}
@@ -1578,10 +1735,10 @@ impl Compiler<'_> {
15781735
let idx = self.name(attr.as_str());
15791736
emit!(self, Instruction::DeleteAttr { idx });
15801737
}
1581-
Expr::Subscript(ExprSubscript { value, slice, .. }) => {
1582-
self.compile_expression(value)?;
1583-
self.compile_expression(slice)?;
1584-
emit!(self, Instruction::DeleteSubscript);
1738+
Expr::Subscript(ExprSubscript {
1739+
value, slice, ctx, ..
1740+
}) => {
1741+
self.compile_subscript(value, slice, *ctx)?;
15851742
}
15861743
Expr::Tuple(ExprTuple { elts, .. }) | Expr::List(ExprList { elts, .. }) => {
15871744
for element in elts {
@@ -3946,10 +4103,10 @@ impl Compiler<'_> {
39464103
fn compile_store(&mut self, target: &Expr) -> CompileResult<()> {
39474104
match &target {
39484105
Expr::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);
4106+
Expr::Subscript(ExprSubscript {
4107+
value, slice, ctx, ..
4108+
}) => {
4109+
self.compile_subscript(value, slice, *ctx)?;
39534110
}
39544111
Expr::Attribute(ExprAttribute { value, attr, .. }) => {
39554112
self.check_forbidden_name(attr.as_str(), NameUsage::Store)?;
@@ -4030,7 +4187,14 @@ impl Compiler<'_> {
40304187
self.compile_name(id, NameUsage::Load)?;
40314188
AugAssignKind::Name { id }
40324189
}
4033-
Expr::Subscript(ExprSubscript { value, slice, .. }) => {
4190+
Expr::Subscript(ExprSubscript {
4191+
value,
4192+
slice,
4193+
ctx: _,
4194+
..
4195+
}) => {
4196+
// For augmented assignment, we need to load the value first
4197+
// But we can't use compile_subscript directly because we need DUP_TOP2
40344198
self.compile_expression(value)?;
40354199
self.compile_expression(slice)?;
40364200
emit!(self, Instruction::Duplicate2);
@@ -4264,10 +4428,10 @@ impl Compiler<'_> {
42644428
// Perform operation:
42654429
self.compile_op(op, false);
42664430
}
4267-
Expr::Subscript(ExprSubscript { value, slice, .. }) => {
4268-
self.compile_expression(value)?;
4269-
self.compile_expression(slice)?;
4270-
emit!(self, Instruction::Subscript);
4431+
Expr::Subscript(ExprSubscript {
4432+
value, slice, ctx, ..
4433+
}) => {
4434+
self.compile_subscript(value, slice, *ctx)?;
42714435
}
42724436
Expr::UnaryOp(ExprUnaryOp { op, operand, .. }) => {
42734437
self.compile_expression(operand)?;
@@ -4298,30 +4462,13 @@ impl Compiler<'_> {
42984462
// self.emit_load_const(compile_constant(value));
42994463
// }
43004464
Expr::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-
}
4465+
self.starunpack_helper(elts, 0, CollectionType::List)?;
43074466
}
43084467
Expr::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-
}
4468+
self.starunpack_helper(elts, 0, CollectionType::Tuple)?;
43174469
}
43184470
Expr::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-
}
4471+
self.starunpack_helper(elts, 0, CollectionType::Set)?;
43254472
}
43264473
Expr::Dict(ExprDict { items, .. }) => {
43274474
self.compile_dict(items)?;

0 commit comments

Comments
 (0)