Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Adjust CodeFlag
  • Loading branch information
ShaharNaveh committed Jan 10, 2026
commit ec25472902234590b13546924677025c64f29a6b
30 changes: 15 additions & 15 deletions crates/codegen/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,7 @@ enum CollectionType {
impl Compiler {
fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self {
let module_code = ir::CodeInfo {
flags: bytecode::CodeFlags::NEW_LOCALS,
flags: bytecode::CodeFlags::NEWLOCALS,
source_path: source_file.name().to_owned(),
private: None,
blocks: vec![ir::Block::default()],
Expand Down Expand Up @@ -749,19 +749,19 @@ impl Compiler {
CompilerScope::Module => (bytecode::CodeFlags::empty(), 0, 0, 0),
CompilerScope::Class => (bytecode::CodeFlags::empty(), 0, 0, 0),
CompilerScope::Function | CompilerScope::AsyncFunction | CompilerScope::Lambda => (
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
bytecode::CodeFlags::NEWLOCALS | bytecode::CodeFlags::OPTIMIZED,
0, // Will be set later in enter_function
0, // Will be set later in enter_function
0, // Will be set later in enter_function
),
CompilerScope::Comprehension => (
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
bytecode::CodeFlags::NEWLOCALS | bytecode::CodeFlags::OPTIMIZED,
0,
1, // comprehensions take one argument (.0)
0,
),
CompilerScope::TypeParams => (
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
bytecode::CodeFlags::NEWLOCALS | bytecode::CodeFlags::OPTIMIZED,
0,
0,
0,
Expand Down Expand Up @@ -1298,7 +1298,7 @@ impl Compiler {
let parent_obj_name = &parent.metadata.name;

// Determine if parent is a function-like scope
let is_function_parent = parent.flags.contains(bytecode::CodeFlags::IS_OPTIMIZED)
let is_function_parent = parent.flags.contains(bytecode::CodeFlags::OPTIMIZED)
&& !parent_obj_name.starts_with("<") // Not a special scope like <lambda>, <listcomp>, etc.
&& parent_obj_name != "<module>"; // Not the module scope

Expand Down Expand Up @@ -1914,7 +1914,7 @@ impl Compiler {
&& self
.current_code_info()
.flags
.contains(bytecode::CodeFlags::IS_GENERATOR)
.contains(bytecode::CodeFlags::GENERATOR)
{
return Err(self.error_ranged(
CodegenErrorType::AsyncReturnValue,
Expand Down Expand Up @@ -2062,7 +2062,7 @@ impl Compiler {
}

self.push_output(
bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
bytecode::CodeFlags::NEWLOCALS | bytecode::CodeFlags::OPTIMIZED,
parameters.posonlyargs.len().to_u32(),
(parameters.posonlyargs.len() + parameters.args.len()).to_u32(),
parameters.kwonlyargs.len().to_u32(),
Expand All @@ -2080,11 +2080,11 @@ impl Compiler {
}

if let Some(name) = parameters.vararg.as_deref() {
self.current_code_info().flags |= bytecode::CodeFlags::HAS_VARARGS;
self.current_code_info().flags |= bytecode::CodeFlags::VARARGS;
self.varname(name.name.as_str())?;
}
if let Some(name) = parameters.kwarg.as_deref() {
self.current_code_info().flags |= bytecode::CodeFlags::HAS_VARKEYWORDS;
self.current_code_info().flags |= bytecode::CodeFlags::VARKEYWORDS;
self.varname(name.name.as_str())?;
}

Expand Down Expand Up @@ -3103,7 +3103,7 @@ impl Compiler {
self.enter_function(name, parameters)?;
self.current_code_info()
.flags
.set(bytecode::CodeFlags::IS_COROUTINE, is_async);
.set(bytecode::CodeFlags::COROUTINE, is_async);

// Set up context
let prev_ctx = self.ctx;
Expand Down Expand Up @@ -3233,7 +3233,7 @@ impl Compiler {
// Enter type params scope
let type_params_name = format!("<generic parameters of {name}>");
self.push_output(
bytecode::CodeFlags::IS_OPTIMIZED | bytecode::CodeFlags::NEW_LOCALS,
bytecode::CodeFlags::OPTIMIZED | bytecode::CodeFlags::NEWLOCALS,
0,
num_typeparam_args as u32,
0,
Expand Down Expand Up @@ -3664,7 +3664,7 @@ impl Compiler {
if is_generic {
let type_params_name = format!("<generic parameters of {name}>");
self.push_output(
bytecode::CodeFlags::IS_OPTIMIZED | bytecode::CodeFlags::NEW_LOCALS,
bytecode::CodeFlags::OPTIMIZED | bytecode::CodeFlags::NEWLOCALS,
0,
0,
0,
Expand Down Expand Up @@ -6361,9 +6361,9 @@ impl Compiler {
in_async_scope: prev_ctx.in_async_scope || is_async,
};

let flags = bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED;
let flags = bytecode::CodeFlags::NEWLOCALS | bytecode::CodeFlags::OPTIMIZED;
let flags = if is_async {
flags | bytecode::CodeFlags::IS_COROUTINE
flags | bytecode::CodeFlags::COROUTINE
} else {
flags
};
Expand Down Expand Up @@ -7030,7 +7030,7 @@ impl Compiler {
}

fn mark_generator(&mut self) {
self.current_code_info().flags |= bytecode::CodeFlags::IS_GENERATOR
self.current_code_info().flags |= bytecode::CodeFlags::GENERATOR
}

/// Whether the expression contains an await expression and
Expand Down
4 changes: 2 additions & 2 deletions crates/codegen/src/ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -283,8 +283,8 @@ impl CodeInfo {

let total_args = self.metadata.argcount
+ self.metadata.kwonlyargcount
+ self.flags.contains(CodeFlags::HAS_VARARGS) as u32
+ self.flags.contains(CodeFlags::HAS_VARKEYWORDS) as u32;
+ self.flags.contains(CodeFlags::VARARGS) as u32
+ self.flags.contains(CodeFlags::VARKEYWORDS) as u32;

let mut found_cellarg = false;
let cell2arg = self
Expand Down
29 changes: 8 additions & 21 deletions crates/compiler-core/src/bytecode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -357,28 +357,15 @@ pub struct CodeObject<C: Constant = ConstantData> {
bitflags! {
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct CodeFlags: u16 {
const NEW_LOCALS = 0x01;
const IS_GENERATOR = 0x02;
const IS_COROUTINE = 0x04;
const HAS_VARARGS = 0x08;
const HAS_VARKEYWORDS = 0x10;
const IS_OPTIMIZED = 0x20;
const OPTIMIZED = 0x0001;
const NEWLOCALS = 0x0002;
const VARARGS = 0x0004;
const VARKEYWORDS = 0x0008;
const GENERATOR = 0x0020;
const COROUTINE = 0x0080;
}
}

impl CodeFlags {
pub const NAME_MAPPING: &'static [(&'static str, Self)] = &[
("GENERATOR", Self::IS_GENERATOR),
("COROUTINE", Self::IS_COROUTINE),
(
"ASYNC_GENERATOR",
Self::from_bits_truncate(Self::IS_GENERATOR.bits() | Self::IS_COROUTINE.bits()),
),
("VARARGS", Self::HAS_VARARGS),
("VARKEYWORDS", Self::HAS_VARKEYWORDS),
];
}

/// an opcode argument that may be extended by a prior ExtendedArg
#[derive(Copy, Clone, PartialEq, Eq)]
#[repr(transparent)]
Expand Down Expand Up @@ -1550,14 +1537,14 @@ impl<C: Constant> CodeObject<C> {
let args = &self.varnames[..nargs];
let kwonlyargs = &self.varnames[nargs..varargs_pos];

let vararg = if self.flags.contains(CodeFlags::HAS_VARARGS) {
let vararg = if self.flags.contains(CodeFlags::VARARGS) {
let vararg = &self.varnames[varargs_pos];
varargs_pos += 1;
Some(vararg)
} else {
None
};
let varkwarg = if self.flags.contains(CodeFlags::HAS_VARKEYWORDS) {
let varkwarg = if self.flags.contains(CodeFlags::VARKEYWORDS) {
Some(&self.varnames[varargs_pos])
} else {
None
Expand Down
10 changes: 5 additions & 5 deletions crates/vm/src/builtins/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ impl PyFunction {

let mut vararg_offset = total_args;
// Pack other positional arguments in to *args:
if code.flags.contains(bytecode::CodeFlags::HAS_VARARGS) {
if code.flags.contains(bytecode::CodeFlags::VARARGS) {
let vararg_value = vm.ctx.new_tuple(args_iter.collect());
fastlocals[vararg_offset] = Some(vararg_value.into());
vararg_offset += 1;
Expand Down Expand Up @@ -155,7 +155,7 @@ impl PyFunction {
}

// Do we support `**kwargs` ?
let kwargs = if code.flags.contains(bytecode::CodeFlags::HAS_VARKEYWORDS) {
let kwargs = if code.flags.contains(bytecode::CodeFlags::VARKEYWORDS) {
let d = vm.ctx.new_dict();
fastlocals[vararg_offset] = Some(d.clone().into());
Some(d)
Expand Down Expand Up @@ -414,7 +414,7 @@ impl Py<PyFunction> {

let code = self.code.lock().clone();

let locals = if code.flags.contains(bytecode::CodeFlags::NEW_LOCALS) {
let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) {
ArgMapping::from_dict_exact(vm.ctx.new_dict())
} else if let Some(locals) = locals {
locals
Expand All @@ -436,8 +436,8 @@ impl Py<PyFunction> {
self.fill_locals_from_args(&frame, func_args, vm)?;

// If we have a generator, create a new generator
let is_gen = code.flags.contains(bytecode::CodeFlags::IS_GENERATOR);
let is_coro = code.flags.contains(bytecode::CodeFlags::IS_COROUTINE);
let is_gen = code.flags.contains(bytecode::CodeFlags::GENERATOR);
let is_coro = code.flags.contains(bytecode::CodeFlags::COROUTINE);
Comment on lines +439 to +440
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

LGTM! Flag renames align with standard conventions.

The renames from IS_GENERATOR to GENERATOR and IS_COROUTINE to COROUTINE are consistent with standard Python CodeFlags naming, removing the IS_ prefix. The match logic correctly handles all combinations for creating the appropriate generator/coroutine/async generator objects.

🤖 Prompt for AI Agents
In @crates/vm/src/builtins/function.rs around lines 439 - 440, No code changes
required; the rename from IS_GENERATOR/IS_COROUTINE to
bytecode::CodeFlags::GENERATOR and ::COROUTINE and the match logic that
constructs generator/coroutine/async generator objects are correct as-is, so
leave the uses of code.flags.contains(bytecode::CodeFlags::GENERATOR) and
code.flags.contains(bytecode::CodeFlags::COROUTINE) and the surrounding
match/creation logic unchanged.

match (is_gen, is_coro) {
(true, false) => {
Ok(PyGenerator::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm))
Expand Down
10 changes: 3 additions & 7 deletions crates/vm/src/frame.rs
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ impl Frame {
Ok(())
};
map_to_dict(&code.cellvars, &self.cells_frees)?;
if code.flags.contains(bytecode::CodeFlags::IS_OPTIMIZED) {
if code.flags.contains(bytecode::CodeFlags::OPTIMIZED) {
map_to_dict(&code.freevars, &self.cells_frees[code.cellvars.len()..])?;
}
}
Expand Down Expand Up @@ -1000,11 +1000,7 @@ impl ExecutingFrame<'_> {
let iterable = self.pop_value();
let iter = if iterable.class().is(vm.ctx.types.coroutine_type) {
// Coroutine requires CO_COROUTINE or CO_ITERABLE_COROUTINE flag
if !self
.code
.flags
.intersects(bytecode::CodeFlags::IS_COROUTINE)
{
if !self.code.flags.intersects(bytecode::CodeFlags::COROUTINE) {
return Err(vm.new_type_error(
"cannot 'yield from' a coroutine object in a non-coroutine generator"
.to_owned(),
Expand Down Expand Up @@ -1662,7 +1658,7 @@ impl ExecutingFrame<'_> {
// arg=0: direct yield (wrapped for async generators)
// arg=1: yield from await/yield-from (NOT wrapped)
let wrap = oparg.get(arg) == 0;
let value = if wrap && self.code.flags.contains(bytecode::CodeFlags::IS_COROUTINE) {
let value = if wrap && self.code.flags.contains(bytecode::CodeFlags::COROUTINE) {
PyAsyncGenWrappedValue(value).into_pyobject(vm)
} else {
value
Expand Down
Loading