diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py index b81115a0c..5415356e5 100644 --- a/Lib/_opcode_metadata.py +++ b/Lib/_opcode_metadata.py @@ -349,6 +349,7 @@ opmap = frozendict( UNPACK_EX=117, UNPACK_SEQUENCE=118, YIELD_VALUE=119, + BUILD_LIST_FROM_ARG=219, INSTRUMENTED_END_FOR=233, INSTRUMENTED_POP_ITER=234, INSTRUMENTED_END_SEND=235, diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 257169e34..1f290d442 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -314,6 +314,10 @@ pub struct CompileOpts { pub dont_imply_dedent: bool, /// Recursion limit used by compiler tree walks, matching Py_EnterRecursiveCall. pub recursion_limit: usize, + /// Emit `BUILD_LIST_FROM_ARG` for a list comprehension with one + /// generator and no `if`, sizing the result from the iterator's length + /// hint. Off by default: the opcode is not part of the CPython set. + pub build_list_from_arg: bool, } impl Default for CompileOpts { @@ -326,6 +330,7 @@ impl Default for CompileOpts { future_features: bytecode::CodeFlags::empty(), dont_imply_dedent: false, recursion_limit: 1000, + build_list_from_arg: false, } } } @@ -10712,6 +10717,24 @@ impl<'warnings> Compiler<'warnings> { Ok(()) } + /// ast.ListComp.build_container_and_load_iter: one generator, no ifs, + /// not an `async for`. The iterator is already on the stack; BUILD_LIST_FROM_ARG + /// sizes the list from its length hint. + fn listcomp_uses_build_list_from_arg( + &self, + init_collection: Option<&AnyInstruction>, + generators: &[ast::Comprehension], + ) -> bool { + self.opts.build_list_from_arg + && matches!( + init_collection, + Some(AnyInstruction::Real(Instruction::BuildList { .. })) + ) + && generators.len() == 1 + && generators[0].ifs.is_empty() + && !generators[0].is_async + } + #[expect(clippy::too_many_arguments, reason = "ignore warning for now")] fn compile_comprehension( &mut self, @@ -10845,8 +10868,15 @@ impl<'warnings> Compiler<'warnings> { None }; + // ast.ListComp.build_container_and_load_iter: one generator and no + // ifs loads the iterator first, then BUILD_LIST_FROM_ARG. + let use_build_list_from_arg = + self.listcomp_uses_build_list_from_arg(init_collection.as_ref(), generators); + // Create empty object of proper type: - if let Some(init_collection) = init_collection { + if let Some(init_collection) = init_collection + && !use_build_list_from_arg + { self._emit(init_collection, OpArg::new(0), BlockIdx::NULL) } @@ -10896,6 +10926,9 @@ impl<'warnings> Compiler<'warnings> { } else { emit!(self, Instruction::GetIter { mode: 0 }); } + if use_build_list_from_arg { + emit!(self, Instruction::BuildListFromArg); + } } } else { // Evaluate iterated item: @@ -11217,8 +11250,14 @@ impl<'warnings> Compiler<'warnings> { None }; - // Step 4: Create the collection (list/set/dict) - if let Some(init_collection) = init_collection { + // Step 4: Create the collection (list/set/dict). + // BUILD_LIST_FROM_ARG pops the iterator, pushes the list, then + // the iterator: the same stack as BUILD_LIST 0; SWAP 2. + let use_build_list_from_arg = + self.listcomp_uses_build_list_from_arg(init_collection.as_ref(), generators); + if use_build_list_from_arg { + emit!(self, Instruction::BuildListFromArg); + } else if let Some(init_collection) = init_collection { self._emit(init_collection, OpArg::new(0), BlockIdx::NULL); // SWAP to get iterator on top emit!(self, Instruction::Swap { i: 2 }); @@ -17552,6 +17591,39 @@ def f(self): ); } + #[test] + fn listcomp_build_list_from_arg_option() { + let has = |source: &str, build_list_from_arg: bool| { + let code = compile_exec_with_options( + source, + CompileOpts { + build_list_from_arg, + ..CompileOpts::default() + }, + ); + let f = find_code(&code, "f").expect("missing f"); + ( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BuildListFromArg)), + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BuildList { .. })), + ) + }; + let single = "def f(n):\n return [i for i in range(n)]\n"; + // (BUILD_LIST_FROM_ARG, BUILD_LIST) + assert_eq!(has(single, false), (false, true)); + assert_eq!(has(single, true), (true, false)); + for source in [ + "def f(xs):\n return [x for x in xs if x]\n", + "def f(xss):\n return [x for xs in xss for x in xs]\n", + "async def f(xs):\n return [x async for x in xs]\n", + ] { + assert_eq!(has(source, true), (false, true), "source={source:?}"); + } + } + #[test] fn or_return_not_taken_before_jump_target_splits_exception_table_like_cpython() { let code = compile_exec_with_options( diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index a34759a92..9d42b5d74 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -680,6 +680,7 @@ define_opcodes!( UnpackSequenceList = 216, UnpackSequenceTuple = 217, UnpackSequenceTwoTuple = 218, + BuildListFromArg = 219, InstrumentedEndFor = 233, InstrumentedPopIter = 234, InstrumentedEndSend = 235, diff --git a/crates/compiler-core/src/bytecode/opcode_metadata.rs b/crates/compiler-core/src/bytecode/opcode_metadata.rs index 82c3ec091..b264639fd 100644 --- a/crates/compiler-core/src/bytecode/opcode_metadata.rs +++ b/crates/compiler-core/src/bytecode/opcode_metadata.rs @@ -822,6 +822,10 @@ impl super::Opcode { Self::UnpackSequenceList => (oparg, 1), Self::UnpackSequenceTuple => (oparg, 1), Self::UnpackSequenceTwoTuple => (2, 1), + Self::BuildListFromArg => ( + 2, // TODO: Differs from CPython `None` + 1, // TODO: Differs from CPython `None` + ), Self::InstrumentedEndFor => (2, 3), Self::InstrumentedPopIter => (0, 2), Self::InstrumentedEndSend => (1, 3), diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 67530e28d..00b2b7e78 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -4585,6 +4585,23 @@ impl ExecutingFrame<'_> { self.push_value(list_obj.into()); Ok(None) } + Instruction::BuildListFromArg => { + let last_val = self.pop_value(); + let length_hint = match last_val.clone().length_hint(0, vm) { + Ok(hint) => hint, + Err(exc) + if exc.fast_isinstance(vm.ctx.exceptions.system_exit) + || exc.fast_isinstance(vm.ctx.exceptions.keyboard_interrupt) => + { + return Err(exc); + } + Err(_) => 0, + }; + let list_obj = vm.ctx.new_list(Vec::with_capacity(length_hint)); + self.push_value(list_obj.into()); + self.push_value(last_val); + Ok(None) + } Instruction::BuildMap { count: size } => self.execute_build_map(vm, size.get(arg)), Instruction::BuildSet { count: size } => { let set = PySet::default().into_ref(&vm.ctx); diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 4d003b8a2..7ae959a3b 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -3119,6 +3119,7 @@ impl VirtualMachine { future_features: crate::bytecode::CodeFlags::empty(), dont_imply_dedent: false, recursion_limit: self.recursion_limit.get(), + build_list_from_arg: false, } } diff --git a/tools/opcode_metadata/conf.toml b/tools/opcode_metadata/conf.toml index e69de29bb..765ecd1b4 100644 --- a/tools/opcode_metadata/conf.toml +++ b/tools/opcode_metadata/conf.toml @@ -0,0 +1,2 @@ +[BuildListFromArg] +stack_effect = { pushed = "2", popped = "1" }