diff --git a/Cargo.lock b/Cargo.lock index 0f0eced17..0a30617fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1534,6 +1534,7 @@ version = "0.1.1" dependencies = [ "bincode", "bitflags", + "itertools", "lz4-compress", "num-bigint", "num-complex", diff --git a/bytecode/Cargo.toml b/bytecode/Cargo.toml index fa72e3451..e3490e316 100644 --- a/bytecode/Cargo.toml +++ b/bytecode/Cargo.toml @@ -15,3 +15,4 @@ lz4-compress = "0.1.1" num-bigint = { version = "0.2", features = ["serde"] } num-complex = { version = "0.2", features = ["serde"] } serde = { version = "1.0", features = ["derive"] } +itertools = "0.8" diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs index 18ae446d9..1db3d3df0 100644 --- a/bytecode/src/bytecode.rs +++ b/bytecode/src/bytecode.rs @@ -2,6 +2,7 @@ //! implements bytecode structure. use bitflags::bitflags; +use itertools::Itertools; use num_bigint::BigInt; use num_complex::Complex64; use serde::{Deserialize, Serialize}; @@ -40,9 +41,9 @@ pub struct CodeObject { pub flags: CodeFlags, pub posonlyarg_count: usize, // Number of positional-only arguments pub arg_names: Vec, // Names of positional arguments - pub varargs: Varargs, // *args or * + pub varargs_name: Option, // *args or * pub kwonlyarg_names: Vec, - pub varkeywords: Varargs, // **kwargs or ** + pub varkeywords_name: Option, // **kwargs or ** pub source_path: String, pub first_line_number: usize, pub obj_name: String, // Name of the object that created this code object @@ -57,6 +58,8 @@ bitflags! { const NEW_LOCALS = 0x08; const IS_GENERATOR = 0x10; const IS_COROUTINE = 0x20; + const HAS_VARARGS = 0x40; + const HAS_VARKEYWORDS = 0x80; } } @@ -70,6 +73,8 @@ impl CodeFlags { pub const NAME_MAPPING: &'static [(&'static str, CodeFlags)] = &[ ("GENERATOR", CodeFlags::IS_GENERATOR), ("COROUTINE", CodeFlags::IS_COROUTINE), + ("VARARGS", CodeFlags::HAS_VARARGS), + ("VARKEYWORDS", CodeFlags::HAS_VARKEYWORDS), ]; } @@ -352,13 +357,6 @@ pub enum UnaryOperator { Plus, } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub enum Varargs { - None, - Unnamed, - Named(String), -} - /* Maintain a stack of blocks on the VM. pub enum BlockType { @@ -373,9 +371,9 @@ impl CodeObject { flags: CodeFlags, posonlyarg_count: usize, arg_names: Vec, - varargs: Varargs, + varargs_name: Option, kwonlyarg_names: Vec, - varkeywords: Varargs, + varkeywords_name: Option, source_path: String, first_line_number: usize, obj_name: String, @@ -387,9 +385,9 @@ impl CodeObject { flags, posonlyarg_count, arg_names, - varargs, + varargs_name, kwonlyarg_names, - varkeywords, + varkeywords_name, source_path, first_line_number, obj_name, @@ -418,6 +416,29 @@ impl CodeObject { }) } + pub fn varnames(&self) -> impl Iterator + '_ { + self.arg_names + .iter() + .map(String::as_str) + .chain(self.kwonlyarg_names.iter().map(String::as_str)) + .chain( + self.instructions + .iter() + .filter_map(|i| match i { + Instruction::LoadName { + name, + scope: NameScope::Local, + } + | Instruction::StoreName { + name, + scope: NameScope::Local, + } => Some(name.as_str()), + _ => None, + }) + .unique(), + ) + } + fn display_inner( &self, f: &mut fmt::Formatter, diff --git a/compiler/src/compile.rs b/compiler/src/compile.rs index 8f88965b7..9267c900d 100644 --- a/compiler/src/compile.rs +++ b/compiler/src/compile.rs @@ -14,7 +14,7 @@ use crate::symboltable::{ }; use itertools::Itertools; use num_complex::Complex64; -use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Label, Varargs}; +use rustpython_bytecode::bytecode::{self, CallType, CodeObject, Instruction, Label}; use rustpython_parser::{ast, parser}; type BasicOutputStream = PeepholeOptimizer; @@ -177,9 +177,9 @@ impl Compiler { Default::default(), 0, Vec::new(), - Varargs::None, + None, Vec::new(), - Varargs::None, + None, self.source_path.clone().unwrap(), line_number, obj_name, @@ -686,14 +686,29 @@ impl Compiler { flags |= bytecode::CodeFlags::HAS_KW_ONLY_DEFAULTS; } + let mut compile_varargs = |va: &ast::Varargs, flag| match va { + ast::Varargs::None => None, + ast::Varargs::Unnamed => { + flags |= flag; + None + } + ast::Varargs::Named(name) => { + flags |= flag; + Some(name.arg.clone()) + } + }; + + let varargs_name = compile_varargs(&args.vararg, bytecode::CodeFlags::HAS_VARARGS); + let varkeywords_name = compile_varargs(&args.kwarg, bytecode::CodeFlags::HAS_VARKEYWORDS); + let line_number = self.get_source_line_number(); self.push_output(CodeObject::new( flags, args.posonlyargs_count, args.args.iter().map(|a| a.arg.clone()).collect(), - compile_varargs(&args.vararg), + varargs_name, args.kwonlyargs.iter().map(|a| a.arg.clone()).collect(), - compile_varargs(&args.kwarg), + varkeywords_name, self.source_path.clone().unwrap(), line_number, name.to_owned(), @@ -970,9 +985,9 @@ impl Compiler { Default::default(), 0, vec![], - Varargs::None, + None, vec![], - Varargs::None, + None, self.source_path.clone().unwrap(), line_number, name.to_owned(), @@ -1920,9 +1935,9 @@ impl Compiler { Default::default(), 1, vec![".0".to_owned()], - Varargs::None, + None, vec![], - Varargs::None, + None, self.source_path.clone().unwrap(), line_number, name.clone(), @@ -2223,14 +2238,6 @@ fn compile_location(location: &ast::Location) -> bytecode::Location { bytecode::Location::new(location.row(), location.column()) } -fn compile_varargs(varargs: &ast::Varargs) -> bytecode::Varargs { - match varargs { - ast::Varargs::None => bytecode::Varargs::None, - ast::Varargs::Unnamed => bytecode::Varargs::Unnamed, - ast::Varargs::Named(param) => bytecode::Varargs::Named(param.arg.clone()), - } -} - fn compile_conversion_flag(conversion_flag: ast::ConversionFlag) -> bytecode::ConversionFlag { match conversion_flag { ast::ConversionFlag::Ascii => bytecode::ConversionFlag::Ascii, diff --git a/vm/src/obj/objcode.rs b/vm/src/obj/objcode.rs index 16410e9be..d46566243 100644 --- a/vm/src/obj/objcode.rs +++ b/vm/src/obj/objcode.rs @@ -106,6 +106,16 @@ impl PyCodeRef { fn co_flags(self) -> u8 { self.code.flags.bits() } + + #[pyproperty] + fn co_varnames(self, vm: &VirtualMachine) -> PyObjectRef { + let varnames = self + .code + .varnames() + .map(|s| vm.new_str(s.to_owned())) + .collect(); + vm.ctx.new_tuple(varnames) + } } pub fn init(ctx: &PyContext) { diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs index 703c8bfae..b485ec7a4 100644 --- a/vm/src/obj/objfunction.rs +++ b/vm/src/obj/objfunction.rs @@ -93,37 +93,37 @@ impl PyFunction { } // Pack other positional arguments in to *args: - match code_object.varargs { - bytecode::Varargs::Named(ref vararg_name) => { - let mut last_args = vec![]; - for i in n..nargs { - let arg = &func_args.args[i]; - last_args.push(arg.clone()); - } - let vararg_value = vm.ctx.new_tuple(last_args); - - locals.set_item(vararg_name, vararg_value, vm)?; + if let Some(ref vararg_name) = code_object.varargs_name { + let mut last_args = vec![]; + for i in n..nargs { + let arg = &func_args.args[i]; + last_args.push(arg.clone()); } - bytecode::Varargs::Unnamed | bytecode::Varargs::None => { - // Check the number of positional arguments - if nargs > nexpected_args { - return Err(vm.new_type_error(format!( - "Expected {} arguments (got: {})", - nexpected_args, nargs - ))); - } + let vararg_value = vm.ctx.new_tuple(last_args); + + locals.set_item(vararg_name, vararg_value, vm)?; + } else { + // Check the number of positional arguments + if nargs > nexpected_args { + return Err(vm.new_type_error(format!( + "Expected {} arguments (got: {})", + nexpected_args, nargs + ))); } } // Do we support `**kwargs` ? - let kwargs = match code_object.varkeywords { - bytecode::Varargs::Named(ref kwargs_name) => { - let d = vm.ctx.new_dict(); + let kwargs = if code_object + .flags + .contains(bytecode::CodeFlags::HAS_VARKEYWORDS) + { + let d = vm.ctx.new_dict(); + if let Some(ref kwargs_name) = code_object.varkeywords_name { locals.set_item(kwargs_name, d.as_object().clone(), vm)?; - Some(d) } - bytecode::Varargs::Unnamed => Some(vm.ctx.new_dict()), - bytecode::Varargs::None => None, + Some(d) + } else { + None }; // Handle keyword arguments