Support StringIdx in the vm

Author: coolreader18

src/shell/helper.rs

@@ -1,5 +1,5 @@
 use rustpython_vm::obj::objstr::PyStringRef;
-use rustpython_vm::pyobject::{PyIterable, PyResult, TryFromObject};
+use rustpython_vm::pyobject::{PyIterable, PyResult, StringRef, TryFromObject};
 use rustpython_vm::scope::{NameProtocol, Scope};
 use rustpython_vm::VirtualMachine;
 
@@ -77,7 +77,9 @@ impl<'vm> ShellHelper<'vm> {
             // last: the last word, could be empty if it ends with a dot
             // parents: the words before the dot
 
-            let mut current = self.scope.load_global(self.vm, first)?;
+            let mut current = self
+                .scope
+                .load_global(self.vm, &StringRef::new(first.into()))?;
 
             for attr in parents {
                 current = self.vm.get_attribute(current.clone(), attr.as_str()).ok()?;

vm/src/dictdatatype.rs

@@ -1,14 +1,12 @@
 use crate::obj::objstr::{PyString, PyStringRef};
 use crate::pyhash;
-use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult};
+use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult, StringRef};
 use crate::vm::VirtualMachine;
-use num_bigint::ToBigInt;
 /// Ordered dictionary implementation.
 /// Inspired by: https://morepypy.blogspot.com/2015/01/faster-more-memory-efficient-and-more.html
 /// And: https://www.youtube.com/watch?v=p33CVV29OG8
 /// And: http://code.activestate.com/recipes/578375/
-use std::collections::{hash_map::DefaultHasher, HashMap};
-use std::hash::{Hash, Hasher};
+use std::collections::HashMap;
 use std::mem::size_of;
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
@@ -423,10 +421,7 @@ pub trait DictKey {
 /// to index dictionaries.
 impl DictKey for &PyObjectRef {
     fn do_hash(self, vm: &VirtualMachine) -> PyResult<HashValue> {
-        let raw_hash = vm._hash(self)?;
-        let mut hasher = DefaultHasher::new();
-        raw_hash.hash(&mut hasher);
-        Ok(hasher.finish() as HashValue)
+        vm._hash(self)
     }
 
     fn do_is(self, other: &PyObjectRef) -> bool {
@@ -438,9 +433,14 @@ impl DictKey for &PyObjectRef {
     }
 }
 
+// for the following DictKey impls, we use pyhash::hash_as_bigint because vm._hash() takes the
+// BigInt that obj.__hash__() returns and does `raw_hash % MODULUS` to get an i64.
+// pyhash::hash_as_bigint does the same here, to ensure that the string types' do_hash returns the
+// same as what PyObjectRef::do_hash() does for an equivalent python str
+
 impl DictKey for &PyStringRef {
     fn do_hash(self, _vm: &VirtualMachine) -> PyResult<HashValue> {
-        Ok(self.hash())
+        Ok(pyhash::hash_as_bigint(self.hash()))
     }
 
     fn do_is(self, other: &PyObjectRef) -> bool {
@@ -458,16 +458,30 @@ impl DictKey for &PyStringRef {
     }
 }
 
+impl DictKey for &StringRef {
+    fn do_hash(self, _vm: &VirtualMachine) -> PyResult<HashValue> {
+        Ok(pyhash::hash_as_bigint(self.hash_value()))
+    }
+
+    fn do_is(self, _other: &PyObjectRef) -> bool {
+        false
+    }
+
+    fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
+        if let Some(py_str_value) = other_key.payload::<PyString>() {
+            Ok(py_str_value.as_str() == self.as_str())
+        } else {
+            vm.bool_eq(vm.new_str(self), other_key.clone())
+        }
+    }
+}
+
 /// Implement trait for the str type, so that we can use strings
 /// to index dictionaries.
 impl DictKey for &str {
     fn do_hash(self, _vm: &VirtualMachine) -> PyResult<HashValue> {
         // follow a similar route as the hashing of PyStringRef
-        let raw_hash = pyhash::hash_value(&self.to_owned()).to_bigint().unwrap();
-        let raw_hash = pyhash::hash_bigint(&raw_hash);
-        let mut hasher = DefaultHasher::new();
-        raw_hash.hash(&mut hasher);
-        Ok(hasher.finish() as HashValue)
+        Ok(pyhash::hash_as_bigint(pyhash::hash_value(self) as HashValue))
     }
 
     fn do_is(self, _other: &PyObjectRef) -> bool {

vm/src/frame.rs

@@ -6,7 +6,7 @@ use indexmap::IndexMap;
 use itertools::Itertools;
 
 use crate::builtins::builtin_isinstance;
-use crate::bytecode;
+use crate::bytecode::{self, StringIdx};
 use crate::exceptions::{self, ExceptionCtor, PyBaseExceptionRef};
 use crate::function::PyFuncArgs;
 use crate::obj::objasyncgenerator::PyAsyncGenWrappedValue;
@@ -348,12 +348,12 @@ impl ExecutingFrame<'_> {
             bytecode::Instruction::LoadName {
                 ref name,
                 ref scope,
-            } => self.load_name(vm, name, scope),
+            } => self.load_name(vm, *name, scope),
             bytecode::Instruction::StoreName {
                 ref name,
                 ref scope,
-            } => self.store_name(vm, name, scope),
-            bytecode::Instruction::DeleteName { ref name } => self.delete_name(vm, name),
+            } => self.store_name(vm, *name, scope),
+            bytecode::Instruction::DeleteName { ref name } => self.delete_name(vm, *name),
             bytecode::Instruction::Subscript => self.execute_subscript(vm),
             bytecode::Instruction::StoreSubscript => self.execute_store_subscript(vm),
             bytecode::Instruction::DeleteSubscript => self.execute_delete_subscript(vm),
@@ -437,9 +437,9 @@ impl ExecutingFrame<'_> {
             bytecode::Instruction::BinaryOperation { ref op, inplace } => {
                 self.execute_binop(vm, op, *inplace)
             }
-            bytecode::Instruction::LoadAttr { ref name } => self.load_attr(vm, name),
-            bytecode::Instruction::StoreAttr { ref name } => self.store_attr(vm, name),
-            bytecode::Instruction::DeleteAttr { ref name } => self.delete_attr(vm, name),
+            bytecode::Instruction::LoadAttr { ref name } => self.load_attr(vm, *name),
+            bytecode::Instruction::StoreAttr { ref name } => self.store_attr(vm, *name),
+            bytecode::Instruction::DeleteAttr { ref name } => self.delete_attr(vm, *name),
             bytecode::Instruction::UnaryOperation { ref op } => self.execute_unop(vm, op),
             bytecode::Instruction::CompareOperation { ref op } => self.execute_compare(vm, op),
             bytecode::Instruction::ReturnValue => {
@@ -786,9 +786,8 @@ impl ExecutingFrame<'_> {
         if let Some(dict) = module.dict() {
             for (k, v) in &dict {
                 let k = vm.to_str(&k)?;
-                let k = k.as_str();
-                if !k.starts_with('_') {
-                    self.scope.store_name(&vm, k, v);
+                if !k.as_str().starts_with('_') {
+                    self.scope.store_name(&vm, &k.data(), v);
                 }
             }
         }
@@ -859,28 +858,22 @@ impl ExecutingFrame<'_> {
     fn store_name(
         &mut self,
         vm: &VirtualMachine,
-        name: &str,
+        name: StringIdx,
         name_scope: &bytecode::NameScope,
     ) -> FrameResult {
         let obj = self.pop_value();
+        let name = self.code.get_string(name);
         match name_scope {
-            bytecode::NameScope::Global => {
-                self.scope.store_global(vm, name, obj);
-            }
-            bytecode::NameScope::NonLocal => {
-                self.scope.store_cell(vm, name, obj);
-            }
-            bytecode::NameScope::Local => {
-                self.scope.store_name(vm, name, obj);
-            }
-            bytecode::NameScope::Free => {
-                self.scope.store_name(vm, name, obj);
-            }
+            bytecode::NameScope::Global => self.scope.store_global(vm, name, obj),
+            bytecode::NameScope::NonLocal => self.scope.store_cell(vm, name, obj),
+            bytecode::NameScope::Local => self.scope.store_name(vm, name, obj),
+            bytecode::NameScope::Free => self.scope.store_name(vm, name, obj),
         }
         Ok(None)
     }
 
-    fn delete_name(&self, vm: &VirtualMachine, name: &str) -> FrameResult {
+    fn delete_name(&self, vm: &VirtualMachine, name: StringIdx) -> FrameResult {
+        let name = self.code.get_string(name);
         match self.scope.delete_name(vm, name) {
             Ok(_) => Ok(None),
             Err(_) => Err(vm.new_name_error(format!("name '{}' is not defined", name))),
@@ -891,9 +884,10 @@ impl ExecutingFrame<'_> {
     fn load_name(
         &mut self,
         vm: &VirtualMachine,
-        name: &str,
+        name: StringIdx,
         name_scope: &bytecode::NameScope,
     ) -> FrameResult {
+        let name = self.code.get_string(name);
         let optional_value = match name_scope {
             bytecode::NameScope::Global => self.scope.load_global(vm, name),
             bytecode::NameScope::NonLocal => self.scope.load_cell(vm, name),
@@ -1412,23 +1406,25 @@ impl ExecutingFrame<'_> {
         Ok(None)
     }
 
-    fn load_attr(&mut self, vm: &VirtualMachine, attr_name: &str) -> FrameResult {
+    fn load_attr(&mut self, vm: &VirtualMachine, attr_name: StringIdx) -> FrameResult {
+        let attr_name = vm.new_str(self.code.get_string(attr_name));
         let parent = self.pop_value();
         let obj = vm.get_attribute(parent, attr_name)?;
         self.push_value(obj);
         Ok(None)
     }
 
-    fn store_attr(&mut self, vm: &VirtualMachine, attr_name: &str) -> FrameResult {
+    fn store_attr(&mut self, vm: &VirtualMachine, attr_name: StringIdx) -> FrameResult {
+        let attr_name = vm.new_str(self.code.get_string(attr_name));
         let parent = self.pop_value();
         let value = self.pop_value();
-        vm.set_attr(&parent, vm.new_str(attr_name.to_owned()), value)?;
+        vm.set_attr(&parent, attr_name, value)?;
         Ok(None)
     }
 
-    fn delete_attr(&mut self, vm: &VirtualMachine, attr_name: &str) -> FrameResult {
+    fn delete_attr(&mut self, vm: &VirtualMachine, attr_name: StringIdx) -> FrameResult {
         let parent = self.pop_value();
-        let name = vm.ctx.new_str(attr_name.to_owned());
+        let name = vm.new_str(self.code.get_string(attr_name));
         vm.del_attr(&parent, name)?;
         Ok(None)
     }

vm/src/lib.rs

@@ -88,5 +88,6 @@ pub use rustpython_bytecode::*;
 #[doc(hidden)]
 pub mod __exports {
     pub use maplit::hashmap;
+    pub use once_cell::sync::Lazy;
     pub use smallbox::smallbox;
 }

vm/src/macros.rs

@@ -342,3 +342,14 @@ macro_rules! named_function {
         }
     }};
 }
+
+#[macro_export]
+macro_rules! strdata {
+    ($s:literal) => {{
+        use $crate::__exports::Lazy;
+        static DATA: Lazy<$crate::pyobject::StringRef> =
+            Lazy::new(|| $crate::pyobject::StringRef::new($s.into()));
+        let data: &$crate::pyobject::StringRef = &DATA;
+        data
+    }};
+}

vm/src/obj/objfunction.rs

@@ -81,21 +81,18 @@ impl PyFunction {
         // See also: PyEval_EvalCodeWithName in cpython:
         // https://github.com/python/cpython/blob/master/Python/ceval.c#L3681
 
-        let n = if nargs > nexpected_args {
-            nexpected_args
-        } else {
-            nargs
-        };
+        let n = std::cmp::min(nargs, nexpected_args);
 
         // Copy positional arguments into local variables
         for i in 0..n {
-            let arg_name = &code_object.arg_names[i];
+            let arg_name = code_object.get_string(code_object.arg_names[i]);
             let arg = &func_args.args[i];
             locals.set_item(arg_name, arg.clone(), vm)?;
         }
 
         // Pack other positional arguments in to *args:
-        if let Some(ref vararg_name) = code_object.varargs_name {
+        if let Some(vararg_name) = code_object.varargs_name {
+            let vararg_name = code_object.get_string(vararg_name);
             let mut last_args = vec![];
             for i in n..nargs {
                 let arg = &func_args.args[i];
@@ -120,7 +117,8 @@ impl PyFunction {
             .contains(bytecode::CodeFlags::HAS_VARKEYWORDS)
         {
             let d = vm.ctx.new_dict();
-            if let Some(ref kwargs_name) = code_object.varkeywords_name {
+            if let Some(kwargs_name) = code_object.varkeywords_name {
+                let kwargs_name = code_object.get_string(kwargs_name);
                 locals.set_item(kwargs_name, d.as_object().clone(), vm)?;
             }
             Some(d)
@@ -132,9 +130,14 @@ impl PyFunction {
         // Handle keyword arguments
         for (name, value) in func_args.kwargs {
             // Check if we have a parameter with this name:
-            if code_object.arg_names.contains(&name) || code_object.kwonlyarg_names.contains(&name)
+            let contains_str = |v: &[crate::bytecode::StringIdx], s| {
+                v.iter()
+                    .any(|&idx| code_object.get_string(idx).as_str() == s)
+            };
+            if contains_str(&code_object.arg_names, &name)
+                || contains_str(&code_object.kwonlyarg_names, &name)
             {
-                if posonly_args.contains(&name) {
+                if contains_str(&posonly_args, &name) {
                     posonly_passed_as_kwarg.push(name);
                     continue;
                 } else if locals.contains_key(&name, vm) {
@@ -170,24 +173,24 @@ impl PyFunction {
             let required_args = nexpected_args - num_defaults_available;
             let mut missing = vec![];
             for i in 0..required_args {
-                let variable_name = &code_object.arg_names[i];
+                let variable_name = code_object.get_string(code_object.arg_names[i]);
                 if !locals.contains_key(variable_name, vm) {
                     missing.push(variable_name)
                 }
             }
             if !missing.is_empty() {
                 return Err(vm.new_type_error(format!(
-                    "Missing {} required positional arguments: {:?}",
+                    "Missing {} required positional arguments: {}",
                     missing.len(),
-                    missing
+                    missing.iter().format(", "),
                 )));
             }
             if let Some(defaults) = &self.defaults {
                 let defaults = defaults.as_slice();
                 // We have sufficient defaults, so iterate over the corresponding names and use
                 // the default if we don't already have a value
                 for (default_index, i) in (required_args..nexpected_args).enumerate() {
-                    let arg_name = &code_object.arg_names[i];
+                    let arg_name = code_object.get_string(code_object.arg_names[i]);
                     if !locals.contains_key(arg_name, vm) {
                         locals.set_item(arg_name, defaults[default_index].clone(), vm)?;
                     }
@@ -197,6 +200,7 @@ impl PyFunction {
 
         // Check if kw only arguments are all present:
         for arg_name in &code_object.kwonlyarg_names {
+            let arg_name = code_object.get_string(*arg_name);
             if !locals.contains_key(arg_name, vm) {
                 if let Some(kw_only_defaults) = &self.kw_only_defaults {
                     if let Some(default) = kw_only_defaults.get_item_option(arg_name, vm)? {