Move varargs information into CodeFlags

RustPython · coolreader18 · Apr 4, 2020 · Apr 2, 2020 · Apr 2, 2020 · Apr 2, 2020
commit 6553bf72d40a0b9b700c2db72f4de64dcaedeafc
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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<String>,  // Names of positional arguments
-    pub varargs: Varargs,        // *args or *
+    pub varargs_name: Option<String>, // *args or *
     pub kwonlyarg_names: Vec<String>,
-    pub varkeywords: Varargs, // **kwargs or **
+    pub varkeywords_name: Option<String>, // **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<String>,
-        varargs: Varargs,
+        varargs_name: Option<String>,
         kwonlyarg_names: Vec<String>,
-        varkeywords: Varargs,
+        varkeywords_name: Option<String>,
         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<Item = &str> + '_ {
+        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
@@ -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<CodeObjectStream>;
@@ -177,9 +177,9 @@ impl<O: OutputStream> Compiler<O> {
             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<O: OutputStream> Compiler<O> {
             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<O: OutputStream> Compiler<O> {
             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<O: OutputStream> Compiler<O> {
             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
@@ -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
@@ -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