Skip to content

align varnames to CPython #6894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
youknowone wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

align varnames to CPython #6894

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ddb951d
align varnames to CPython
youknowone Jan 29, 2026
94b06a4
Auto-format: cargo fmt --all
github-actions[bot] Jan 29, 2026
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
29 changes: 22 additions & 7 deletions crates/codegen/src/compile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1015,16 +1015,31 @@ impl Compiler {
// Use varnames from symbol table (already collected in definition order)
let varname_cache: IndexSet<String> = ste.varnames.iter().cloned().collect();

// Build cellvars using dictbytype (CELL scope, sorted)
// Build cellvars in localsplus order:
// 1. cell+local vars in varnames definition order
// 2. cell-only vars sorted alphabetically
let mut cellvar_cache = IndexSet::default();
let mut cell_names: Vec<_> = ste
let cell_scope_names: IndexSet<String> = ste
.symbols
.iter()
.filter(|(_, s)| s.scope == SymbolScope::Cell)
.map(|(name, _)| name.clone())
.collect();
cell_names.sort();
for name in cell_names {

// First: cell vars that are also in varnames (in varnames order)
for var in varname_cache.iter() {
if cell_scope_names.contains(var) {
cellvar_cache.insert(var.clone());
}
}
// Second: cell-only vars (not in varnames, sorted for determinism)
let mut cell_only: Vec<_> = cell_scope_names
.iter()
.filter(|name| !varname_cache.contains(name.as_str()))
.cloned()
.collect();
cell_only.sort();
for name in cell_only {
cellvar_cache.insert(name);
}

Expand Down Expand Up @@ -4168,11 +4183,11 @@ impl Compiler {
flags: bytecode::MakeFunctionFlags,
) -> CompileResult<()> {
// Handle free variables (closure)
let has_freevars = !code.freevars.is_empty();
let has_freevars = !code.freevars().is_empty();
if has_freevars {
// Build closure tuple by loading free variables

for var in &code.freevars {
for var in code.freevars() {
// Special case: If a class contains a method with a
// free variable that has the same name as a method,
// the name will be considered free *and* local in the
Expand Down Expand Up @@ -4229,7 +4244,7 @@ impl Compiler {
emit!(
self,
Instruction::BuildTuple {
size: code.freevars.len().to_u32(),
size: code.freevars().len().to_u32(),
}
);
}
Expand Down
78 changes: 39 additions & 39 deletions crates/codegen/src/ir.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ use rustpython_compiler_core::{
OneIndexed, SourceLocation,
bytecode::{
AnyInstruction, Arg, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData,
ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg,
PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table,
ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label,
LocalKind, OpArg, PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table,
},
varint::{write_signed_varint, write_varint},
};
Expand Down Expand Up @@ -193,7 +193,6 @@ impl CodeInfo {
self.optimize_load_fast_borrow();

let max_stackdepth = self.max_stackdepth()?;
let cell2arg = self.cell2arg();

let Self {
flags,
Expand All @@ -219,13 +218,44 @@ impl CodeInfo {
varnames: varname_cache,
cellvars: cellvar_cache,
freevars: freevar_cache,
fast_hidden: _,
fast_hidden,
argcount: arg_count,
posonlyargcount: posonlyarg_count,
kwonlyargcount: kwonlyarg_count,
firstlineno: first_line_number,
} = metadata;

// Build localsplusnames and localspluskinds
let mut localsplusnames_vec: Vec<String> = Vec::new();
let mut localspluskinds_vec: Vec<LocalKind> = Vec::new();

// 1. For each var in varnames
for var in varname_cache.iter() {
let mut kind = LocalKind::LOCAL;
if cellvar_cache.contains(var) {
kind |= LocalKind::CELL;
}
if fast_hidden.get(var).copied().unwrap_or(false) {
kind |= LocalKind::HIDDEN;
}
localsplusnames_vec.push(var.clone());
localspluskinds_vec.push(kind);
}

// 2. For each var in cellvars that is NOT in varnames
for var in cellvar_cache.iter() {
if !varname_cache.contains(var) {
localsplusnames_vec.push(var.clone());
localspluskinds_vec.push(LocalKind::CELL);
}
}

// 3. For each var in freevars
for var in freevar_cache.iter() {
localsplusnames_vec.push(var.clone());
localspluskinds_vec.push(LocalKind::FREE);
}

let mut instructions = Vec::new();
let mut locations = Vec::new();
let mut linetable_locations: Vec<LineTableLocation> = Vec::new();
Expand Down Expand Up @@ -389,46 +419,16 @@ impl CodeInfo {
locations: locations.into_boxed_slice(),
constants: constants.into_iter().collect(),
names: name_cache.into_iter().collect(),
varnames: varname_cache.into_iter().collect(),
cellvars: cellvar_cache.into_iter().collect(),
freevars: freevar_cache.into_iter().collect(),
cell2arg,
nlocals: varname_cache.len() as u32,
ncellvars: cellvar_cache.len() as u32,
nfreevars: freevar_cache.len() as u32,
localsplusnames: localsplusnames_vec.into_iter().collect(),
localspluskinds: localspluskinds_vec.into_boxed_slice(),
linetable,
exceptiontable,
})
}

fn cell2arg(&self) -> Option<Box<[i32]>> {
if self.metadata.cellvars.is_empty() {
return None;
}

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

let mut found_cellarg = false;
let cell2arg = self
.metadata
.cellvars
.iter()
.map(|var| {
self.metadata
.varnames
.get_index_of(var)
// check that it's actually an arg
.filter(|i| *i < total_args as usize)
.map_or(-1, |i| {
found_cellarg = true;
i as i32
})
})
.collect::<Box<[_]>>();

if found_cellarg { Some(cell2arg) } else { None }
}

fn dce(&mut self) {
for block in &mut self.blocks {
let mut last_instr = None;
Expand Down
80 changes: 58 additions & 22 deletions crates/compiler-core/src/bytecode.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,12 +276,13 @@ pub struct CodeObject<C: Constant = ConstantData> {
pub obj_name: C::Name,
/// Qualified name of the object (like CPython's co_qualname)
pub qualname: C::Name,
pub cell2arg: Option<Box<[i32]>>,
pub localsplusnames: Box<[C::Name]>,
pub localspluskinds: Box<[LocalKind]>,
pub nlocals: u32,
pub ncellvars: u32,
pub nfreevars: u32,
pub constants: Box<[C]>,
pub names: Box<[C::Name]>,
pub varnames: Box<[C::Name]>,
pub cellvars: Box<[C::Name]>,
pub freevars: Box<[C::Name]>,
Comment on lines -282 to -284
Copy link
Member Author

@youknowone youknowone Jan 29, 2026

Choose a reason for hiding this comment

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

@coolreader18 I remember you decided to keep this in current design before. I'd like to listen what do you think about the changes.
The motivation is aligning architecture more to CPython to reduce AI engineering cost. Because no actual technical advantage exists, I will keep our own way if you have good reasons not to do.

/// Line number table (CPython 3.11+ format)
pub linetable: Box<[u8]>,
/// Exception handling table
Expand All @@ -304,6 +305,16 @@ bitflags! {
}
}

bitflags! {
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct LocalKind: u8 {
const LOCAL = 0x20;
const CELL = 0x40;
const FREE = 0x80;
const HIDDEN = 0x10;
}
}

#[derive(Copy, Clone)]
#[repr(C)]
pub struct CodeUnit {
Expand Down Expand Up @@ -561,25 +572,42 @@ impl<N: AsRef<str>> fmt::Debug for Arguments<'_, N> {
}

impl<C: Constant> CodeObject<C> {
/// Returns varnames (first nlocals entries of localsplusnames)
pub fn varnames(&self) -> &[C::Name] {
&self.localsplusnames[..self.nlocals as usize]
}

/// Returns freevars (last nfreevars entries of localsplusnames)
pub fn freevars(&self) -> &[C::Name] {
let start = self.localsplusnames.len() - self.nfreevars as usize;
&self.localsplusnames[start..]
}

/// Total localsplus count
pub fn nlocalsplus(&self) -> usize {
self.localsplusnames.len()
}

/// Get all arguments of the code object
/// like inspect.getargs
pub fn arg_names(&self) -> Arguments<'_, C::Name> {
let varnames = self.varnames();
let nargs = self.arg_count as usize;
let nkwargs = self.kwonlyarg_count as usize;
let mut varargs_pos = nargs + nkwargs;
let posonlyargs = &self.varnames[..self.posonlyarg_count as usize];
let args = &self.varnames[..nargs];
let kwonlyargs = &self.varnames[nargs..varargs_pos];
let posonlyargs = &varnames[..self.posonlyarg_count as usize];
let args = &varnames[..nargs];
let kwonlyargs = &varnames[nargs..varargs_pos];

let vararg = if self.flags.contains(CodeFlags::VARARGS) {
let vararg = &self.varnames[varargs_pos];
let vararg = &varnames[varargs_pos];
varargs_pos += 1;
Some(vararg)
} else {
None
};
let varkwarg = if self.flags.contains(CodeFlags::VARKEYWORDS) {
Some(&self.varnames[varargs_pos])
Some(&varnames[varargs_pos])
} else {
None
};
Expand Down Expand Up @@ -682,9 +710,11 @@ impl<C: Constant> CodeObject<C> {
.map(|x| bag.make_constant(x.borrow_constant()))
.collect(),
names: map_names(self.names),
varnames: map_names(self.varnames),
cellvars: map_names(self.cellvars),
freevars: map_names(self.freevars),
localsplusnames: map_names(self.localsplusnames),
localspluskinds: self.localspluskinds,
nlocals: self.nlocals,
ncellvars: self.ncellvars,
nfreevars: self.nfreevars,
source_path: bag.make_name(self.source_path.as_ref()),
obj_name: bag.make_name(self.obj_name.as_ref()),
qualname: bag.make_name(self.qualname.as_ref()),
Expand All @@ -697,7 +727,6 @@ impl<C: Constant> CodeObject<C> {
kwonlyarg_count: self.kwonlyarg_count,
first_line_number: self.first_line_number,
max_stackdepth: self.max_stackdepth,
cell2arg: self.cell2arg,
linetable: self.linetable,
exceptiontable: self.exceptiontable,
}
Expand All @@ -714,9 +743,11 @@ impl<C: Constant> CodeObject<C> {
.map(|x| bag.make_constant(x.borrow_constant()))
.collect(),
names: map_names(&self.names),
varnames: map_names(&self.varnames),
cellvars: map_names(&self.cellvars),
freevars: map_names(&self.freevars),
localsplusnames: map_names(&self.localsplusnames),
localspluskinds: self.localspluskinds.clone(),
nlocals: self.nlocals,
ncellvars: self.ncellvars,
nfreevars: self.nfreevars,
source_path: bag.make_name(self.source_path.as_ref()),
obj_name: bag.make_name(self.obj_name.as_ref()),
qualname: bag.make_name(self.qualname.as_ref()),
Expand All @@ -729,7 +760,6 @@ impl<C: Constant> CodeObject<C> {
kwonlyarg_count: self.kwonlyarg_count,
first_line_number: self.first_line_number,
max_stackdepth: self.max_stackdepth,
cell2arg: self.cell2arg.clone(),
linetable: self.linetable.clone(),
exceptiontable: self.exceptiontable.clone(),
}
Expand Down Expand Up @@ -773,14 +803,20 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
}

fn get_varname(&self, i: usize) -> &str {
self.varnames[i].as_ref()
self.localsplusnames[i].as_ref()
}

fn get_cell_name(&self, i: usize) -> &str {
self.cellvars
.get(i)
.unwrap_or_else(|| &self.freevars[i - self.cellvars.len()])
.as_ref()
let mut count = 0;
for (name, &kind) in self.localsplusnames.iter().zip(self.localspluskinds.iter()) {
if kind.intersects(LocalKind::CELL | LocalKind::FREE) {
if count == i {
return name.as_ref();
}
count += 1;
}
}
panic!("cell/free index {i} out of bounds")
}
}

Expand Down
Loading
Loading