Skip to content

Commit 6553bf7

Browse files
committed
Move varargs information into CodeFlags
1 parent c1c9631 commit 6553bf7

File tree

6 files changed

+94
-54
lines changed

6 files changed

+94
-54
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

bytecode/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@ lz4-compress = "0.1.1"
1515
num-bigint = { version = "0.2", features = ["serde"] }
1616
num-complex = { version = "0.2", features = ["serde"] }
1717
serde = { version = "1.0", features = ["derive"] }
18+
itertools = "0.8"

bytecode/src/bytecode.rs

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
//! implements bytecode structure.
33
44
use bitflags::bitflags;
5+
use itertools::Itertools;
56
use num_bigint::BigInt;
67
use num_complex::Complex64;
78
use serde::{Deserialize, Serialize};
@@ -40,9 +41,9 @@ pub struct CodeObject {
4041
pub flags: CodeFlags,
4142
pub posonlyarg_count: usize, // Number of positional-only arguments
4243
pub arg_names: Vec<String>, // Names of positional arguments
43-
pub varargs: Varargs, // *args or *
44+
pub varargs_name: Option<String>, // *args or *
4445
pub kwonlyarg_names: Vec<String>,
45-
pub varkeywords: Varargs, // **kwargs or **
46+
pub varkeywords_name: Option<String>, // **kwargs or **
4647
pub source_path: String,
4748
pub first_line_number: usize,
4849
pub obj_name: String, // Name of the object that created this code object
@@ -57,6 +58,8 @@ bitflags! {
5758
const NEW_LOCALS = 0x08;
5859
const IS_GENERATOR = 0x10;
5960
const IS_COROUTINE = 0x20;
61+
const HAS_VARARGS = 0x40;
62+
const HAS_VARKEYWORDS = 0x80;
6063
}
6164
}
6265

@@ -70,6 +73,8 @@ impl CodeFlags {
7073
pub const NAME_MAPPING: &'static [(&'static str, CodeFlags)] = &[
7174
("GENERATOR", CodeFlags::IS_GENERATOR),
7275
("COROUTINE", CodeFlags::IS_COROUTINE),
76+
("VARARGS", CodeFlags::HAS_VARARGS),
77+
("VARKEYWORDS", CodeFlags::HAS_VARKEYWORDS),
7378
];
7479
}
7580

@@ -352,13 +357,6 @@ pub enum UnaryOperator {
352357
Plus,
353358
}
354359

355-
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
356-
pub enum Varargs {
357-
None,
358-
Unnamed,
359-
Named(String),
360-
}
361-
362360
/*
363361
Maintain a stack of blocks on the VM.
364362
pub enum BlockType {
@@ -373,9 +371,9 @@ impl CodeObject {
373371
flags: CodeFlags,
374372
posonlyarg_count: usize,
375373
arg_names: Vec<String>,
376-
varargs: Varargs,
374+
varargs_name: Option<String>,
377375
kwonlyarg_names: Vec<String>,
378-
varkeywords: Varargs,
376+
varkeywords_name: Option<String>,
379377
source_path: String,
380378
first_line_number: usize,
381379
obj_name: String,
@@ -387,9 +385,9 @@ impl CodeObject {
387385
flags,
388386
posonlyarg_count,
389387
arg_names,
390-
varargs,
388+
varargs_name,
391389
kwonlyarg_names,
392-
varkeywords,
390+
varkeywords_name,
393391
source_path,
394392
first_line_number,
395393
obj_name,
@@ -418,6 +416,29 @@ impl CodeObject {
418416
})
419417
}
420418

419+
pub fn varnames(&self) -> impl Iterator<Item = &str> + '_ {
420+
self.arg_names
421+
.iter()
422+
.map(String::as_str)
423+
.chain(self.kwonlyarg_names.iter().map(String::as_str))
424+
.chain(
425+
self.instructions
426+
.iter()
427+
.filter_map(|i| match i {
428+
Instruction::LoadName {
429+
name,
430+
scope: NameScope::Local,
431+
}
432+
| Instruction::StoreName {
433+
name,
434+
scope: NameScope::Local,
435+
} => Some(name.as_str()),
436+
_ => None,
437+
})
438+
.unique(),
439+
)
440+
}
441+
421442
fn display_inner(
422443
&self,
423444
f: &mut fmt::Formatter,

compiler/src/compile.rs

Lines changed: 24 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::symboltable::{
1414
};
1515
use itertools::Itertools;
1616
use num_complex::Complex64;
17-
use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Label, Varargs};
17+
use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Label};
1818
use rustpython_parser::{ast, parser};
1919

2020
type BasicOutputStream = PeepholeOptimizer<CodeObjectStream>;
@@ -177,9 +177,9 @@ impl<O: OutputStream> Compiler<O> {
177177
Default::default(),
178178
0,
179179
Vec::new(),
180-
Varargs::None,
180+
None,
181181
Vec::new(),
182-
Varargs::None,
182+
None,
183183
self.source_path.clone().unwrap(),
184184
line_number,
185185
obj_name,
@@ -686,14 +686,29 @@ impl<O: OutputStream> Compiler<O> {
686686
flags |= bytecode::CodeFlags::HAS_KW_ONLY_DEFAULTS;
687687
}
688688

689+
let mut compile_varargs = |va: &ast::Varargs, flag| match va {
690+
ast::Varargs::None => None,
691+
ast::Varargs::Unnamed => {
692+
flags |= flag;
693+
None
694+
}
695+
ast::Varargs::Named(name) => {
696+
flags |= flag;
697+
Some(name.arg.clone())
698+
}
699+
};
700+
701+
let varargs_name = compile_varargs(&args.vararg, bytecode::CodeFlags::HAS_VARARGS);
702+
let varkeywords_name = compile_varargs(&args.kwarg, bytecode::CodeFlags::HAS_VARKEYWORDS);
703+
689704
let line_number = self.get_source_line_number();
690705
self.push_output(CodeObject::new(
691706
flags,
692707
args.posonlyargs_count,
693708
args.args.iter().map(|a| a.arg.clone()).collect(),
694-
compile_varargs(&args.vararg),
709+
varargs_name,
695710
args.kwonlyargs.iter().map(|a| a.arg.clone()).collect(),
696-
compile_varargs(&args.kwarg),
711+
varkeywords_name,
697712
self.source_path.clone().unwrap(),
698713
line_number,
699714
name.to_owned(),
@@ -970,9 +985,9 @@ impl<O: OutputStream> Compiler<O> {
970985
Default::default(),
971986
0,
972987
vec![],
973-
Varargs::None,
988+
None,
974989
vec![],
975-
Varargs::None,
990+
None,
976991
self.source_path.clone().unwrap(),
977992
line_number,
978993
name.to_owned(),
@@ -1920,9 +1935,9 @@ impl<O: OutputStream> Compiler<O> {
19201935
Default::default(),
19211936
1,
19221937
vec![".0".to_owned()],
1923-
Varargs::None,
1938+
None,
19241939
vec![],
1925-
Varargs::None,
1940+
None,
19261941
self.source_path.clone().unwrap(),
19271942
line_number,
19281943
name.clone(),
@@ -2223,14 +2238,6 @@ fn compile_location(location: &ast::Location) -> bytecode::Location {
22232238
bytecode::Location::new(location.row(), location.column())
22242239
}
22252240

2226-
fn compile_varargs(varargs: &ast::Varargs) -> bytecode::Varargs {
2227-
match varargs {
2228-
ast::Varargs::None => bytecode::Varargs::None,
2229-
ast::Varargs::Unnamed => bytecode::Varargs::Unnamed,
2230-
ast::Varargs::Named(param) => bytecode::Varargs::Named(param.arg.clone()),
2231-
}
2232-
}
2233-
22342241
fn compile_conversion_flag(conversion_flag: ast::ConversionFlag) -> bytecode::ConversionFlag {
22352242
match conversion_flag {
22362243
ast::ConversionFlag::Ascii => bytecode::ConversionFlag::Ascii,

vm/src/obj/objcode.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,16 @@ impl PyCodeRef {
106106
fn co_flags(self) -> u8 {
107107
self.code.flags.bits()
108108
}
109+
110+
#[pyproperty]
111+
fn co_varnames(self, vm: &VirtualMachine) -> PyObjectRef {
112+
let varnames = self
113+
.code
114+
.varnames()
115+
.map(|s| vm.new_str(s.to_owned()))
116+
.collect();
117+
vm.ctx.new_tuple(varnames)
118+
}
109119
}
110120

111121
pub fn init(ctx: &PyContext) {

vm/src/obj/objfunction.rs

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -93,37 +93,37 @@ impl PyFunction {
9393
}
9494

9595
// Pack other positional arguments in to *args:
96-
match code_object.varargs {
97-
bytecode::Varargs::Named(ref vararg_name) => {
98-
let mut last_args = vec![];
99-
for i in n..nargs {
100-
let arg = &func_args.args[i];
101-
last_args.push(arg.clone());
102-
}
103-
let vararg_value = vm.ctx.new_tuple(last_args);
104-
105-
locals.set_item(vararg_name, vararg_value, vm)?;
96+
if let Some(ref vararg_name) = code_object.varargs_name {
97+
let mut last_args = vec![];
98+
for i in n..nargs {
99+
let arg = &func_args.args[i];
100+
last_args.push(arg.clone());
106101
}
107-
bytecode::Varargs::Unnamed | bytecode::Varargs::None => {
108-
// Check the number of positional arguments
109-
if nargs > nexpected_args {
110-
return Err(vm.new_type_error(format!(
111-
"Expected {} arguments (got: {})",
112-
nexpected_args, nargs
113-
)));
114-
}
102+
let vararg_value = vm.ctx.new_tuple(last_args);
103+
104+
locals.set_item(vararg_name, vararg_value, vm)?;
105+
} else {
106+
// Check the number of positional arguments
107+
if nargs > nexpected_args {
108+
return Err(vm.new_type_error(format!(
109+
"Expected {} arguments (got: {})",
110+
nexpected_args, nargs
111+
)));
115112
}
116113
}
117114

118115
// Do we support `**kwargs` ?
119-
let kwargs = match code_object.varkeywords {
120-
bytecode::Varargs::Named(ref kwargs_name) => {
121-
let d = vm.ctx.new_dict();
116+
let kwargs = if code_object
117+
.flags
118+
.contains(bytecode::CodeFlags::HAS_VARKEYWORDS)
119+
{
120+
let d = vm.ctx.new_dict();
121+
if let Some(ref kwargs_name) = code_object.varkeywords_name {
122122
locals.set_item(kwargs_name, d.as_object().clone(), vm)?;
123-
Some(d)
124123
}
125-
bytecode::Varargs::Unnamed => Some(vm.ctx.new_dict()),
126-
bytecode::Varargs::None => None,
124+
Some(d)
125+
} else {
126+
None
127127
};
128128

129129
// Handle keyword arguments

0 commit comments

Comments
 (0)