Skip to content

Commit 5efbe66

Browse files
youknowoneyouknowone
committed
star_unpack_helper
1 parent 7801ef4 commit 5efbe66

File tree

2 files changed

+180
-37
lines changed

2 files changed

+180
-37
lines changed

‎compiler/codegen/src/compile.rs‎

Lines changed: 179 additions & 36 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 subscripter
9+
810
#![deny(clippy::cast_possible_truncation)]
911

1012
use crate::{
@@ -47,9 +49,9 @@ 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-
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,
5355
MatchCase, ModExpression, ModModule, Operator, Parameters, Pattern, PatternMatchAs,
5456
PatternMatchClass, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton,
5557
PatternMatchStar, PatternMatchValue, Singleton, Stmt, StmtExpr, TypeParam, TypeParamParamSpec,
@@ -372,7 +374,158 @@ 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(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+
}
376529
fn error(&mut self, error: CodegenErrorType) -> CodegenError {
377530
self.error_ranged(error, self.current_source_range)
378531
}
@@ -1578,10 +1731,10 @@ impl Compiler<'_> {
15781731
let idx = self.name(attr.as_str());
15791732
emit!(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
}
15861739
Expr::Tuple(ExprTuple { elts, .. }) | Expr::List(ExprList { elts, .. }) => {
15871740
for element in elts {
@@ -3946,10 +4099,10 @@ impl Compiler<'_> {
39464099
fn compile_store(&mut self, target: &Expr) -> CompileResult<()> {
39474100
match &target {
39484101
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);
4102+
Expr::Subscript(ExprSubscript {
4103+
value, slice, ctx, ..
4104+
}) => {
4105+
self.compile_subscript(value, slice, *ctx)?;
39534106
}
39544107
Expr::Attribute(ExprAttribute { value, attr, .. }) => {
39554108
self.check_forbidden_name(attr.as_str(), NameUsage::Store)?;
@@ -4030,7 +4183,14 @@ impl Compiler<'_> {
40304183
self.compile_name(id, NameUsage::Load)?;
40314184
AugAssignKind::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
40344194
self.compile_expression(value)?;
40354195
self.compile_expression(slice)?;
40364196
emit!(self, Instruction::Duplicate2);
@@ -4264,10 +4424,10 @@ impl Compiler<'_> {
42644424
// Perform operation:
42654425
self.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
}
42724432
Expr::UnaryOp(ExprUnaryOp { op, operand, .. }) => {
42734433
self.compile_expression(operand)?;
@@ -4298,30 +4458,13 @@ impl Compiler<'_> {
42984458
// self.emit_load_const(compile_constant(value));
42994459
// }
43004460
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-
}
4461+
self.starunpack_helper(elts, 0, CollectionType::List)?;
43074462
}
43084463
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-
}
4464+
self.starunpack_helper(elts, 0, CollectionType::Tuple)?;
43174465
}
43184466
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-
}
4467+
self.starunpack_helper(elts, 0, CollectionType::Set)?;
43254468
}
43264469
Expr::Dict(ExprDict { items, .. }) => {
43274470
self.compile_dict(items)?;

‎vm/src/frame.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2267,7 +2267,7 @@ impl ExecutingFrame<'_> {
22672267
let list = arg
22682268
.downcast::<PyList>()
22692269
.map_err(|_| vm.new_type_error("LIST_TO_TUPLE expects a list"))?;
2270-
Ok(vm.ctx.new_tuple(list.borrow_vec_mut().clone()).into())
2270+
Ok(vm.ctx.new_tuple(list.borrow_vec().to_vec()).into())
22712271
}
22722272
}
22732273
}

0 commit comments

Comments
 (0)