Allow &str to be used a dictionary key.

src/main.rs

@@ -441,10 +441,10 @@ fn shell_exec(vm: &VirtualMachine, source: &str, scope: Scope) -> Result<(), Com
                 Ok(value) => {
                     // Save non-None values as "_"
 
-                    use rustpython_vm::pyobject::{IdProtocol, IntoPyObject};
+                    use rustpython_vm::pyobject::IdProtocol;
 
                     if !value.is(&vm.get_none()) {
-                        let key = objstr::PyString::from("_").into_pyobject(vm);
+                        let key = "_";
                         scope.globals.set_item(key, value, vm).unwrap();
                     }
                 }

vm/src/dictdatatype.rs

@@ -1,7 +1,7 @@
 use crate::obj::objbool;
 use crate::obj::objstr::PyString;
 use crate::pyhash;
-use crate::pyobject::{IdProtocol, PyObjectRef, PyResult};
+use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult};
 use crate::vm::VirtualMachine;
 use num_bigint::ToBigInt;
 /// Ordered dictionary implementation.
@@ -53,12 +53,12 @@ impl<T: Clone> Dict<T> {
         &mut self,
         hash_index: HashIndex,
         hash_value: HashValue,
-        key: &PyObjectRef,
+        key: PyObjectRef,
         value: T,
     ) {
         let entry = DictEntry {
             hash: hash_value,
-            key: key.clone(),
+            key,
             value,
         };
         let entry_index = self.entries.len();
@@ -73,7 +73,12 @@ impl<T: Clone> Dict<T> {
     }
 
     /// Store a key
-    pub fn insert(&mut self, vm: &VirtualMachine, key: &PyObjectRef, value: T) -> PyResult<()> {
+    pub fn insert<K: DictKey + IntoPyObject + Copy>(
+        &mut self,
+        vm: &VirtualMachine,
+        key: K,
+        value: T,
+    ) -> PyResult<()> {
         match self.lookup(vm, key)? {
             LookupResult::Existing(index) => {
                 // Update existing key
@@ -89,13 +94,13 @@ impl<T: Clone> Dict<T> {
                 hash_value,
             } => {
                 // New key:
-                self.unchecked_push(hash_index, hash_value, key, value);
+                self.unchecked_push(hash_index, hash_value, key.into_pyobject(vm)?, value);
                 Ok(())
             }
         }
     }
 
-    pub fn contains<K: DictKey>(&self, vm: &VirtualMachine, key: &K) -> PyResult<bool> {
+    pub fn contains<K: DictKey + Copy>(&self, vm: &VirtualMachine, key: K) -> PyResult<bool> {
         if let LookupResult::Existing(_) = self.lookup(vm, key)? {
             Ok(true)
         } else {
@@ -113,7 +118,7 @@ impl<T: Clone> Dict<T> {
 
     /// Retrieve a key
     #[cfg_attr(feature = "flame-it", flame("Dict"))]
-    pub fn get<K: DictKey>(&self, vm: &VirtualMachine, key: &K) -> PyResult<Option<T>> {
+    pub fn get<K: DictKey + Copy>(&self, vm: &VirtualMachine, key: K) -> PyResult<Option<T>> {
         if let LookupResult::Existing(index) = self.lookup(vm, key)? {
             Ok(Some(self.unchecked_get(index)))
         } else {
@@ -156,7 +161,7 @@ impl<T: Clone> Dict<T> {
             LookupResult::NewIndex {
                 hash_value,
                 hash_index,
-            } => self.unchecked_push(hash_index, hash_value, &key, value),
+            } => self.unchecked_push(hash_index, hash_value, key.clone(), value),
         };
         Ok(())
     }
@@ -201,7 +206,7 @@ impl<T: Clone> Dict<T> {
 
     /// Lookup the index for the given key.
     #[cfg_attr(feature = "flame-it", flame("Dict"))]
-    fn lookup<K: DictKey>(&self, vm: &VirtualMachine, key: &K) -> PyResult<LookupResult> {
+    fn lookup<K: DictKey + Copy>(&self, vm: &VirtualMachine, key: K) -> PyResult<LookupResult> {
         let hash_value = key.do_hash(vm)?;
         let perturb = hash_value;
         let mut hash_index: HashIndex = hash_value;
@@ -244,7 +249,7 @@ impl<T: Clone> Dict<T> {
     }
 
     /// Retrieve and delete a key
-    pub fn pop<K: DictKey>(&mut self, vm: &VirtualMachine, key: &K) -> PyResult<Option<T>> {
+    pub fn pop<K: DictKey + Copy>(&mut self, vm: &VirtualMachine, key: K) -> PyResult<Option<T>> {
         if let LookupResult::Existing(index) = self.lookup(vm, key)? {
             let value = self.unchecked_get(index);
             self.unchecked_delete(index);
@@ -280,26 +285,26 @@ enum LookupResult {
 /// - PyObjectRef -> arbitrary python type used as key
 /// - str -> string reference used as key, this is often used internally
 pub trait DictKey {
-    fn do_hash(&self, vm: &VirtualMachine) -> PyResult<HashValue>;
-    fn do_is(&self, other: &PyObjectRef) -> bool;
-    fn do_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool>;
+    fn do_hash(self, vm: &VirtualMachine) -> PyResult<HashValue>;
+    fn do_is(self, other: &PyObjectRef) -> bool;
+    fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool>;
 }
 
 /// Implement trait for PyObjectRef such that we can use python objects
 /// to index dictionaries.
-impl DictKey for PyObjectRef {
-    fn do_hash(&self, vm: &VirtualMachine) -> PyResult<HashValue> {
+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)
     }
 
-    fn do_is(&self, other: &PyObjectRef) -> bool {
+    fn do_is(self, other: &PyObjectRef) -> bool {
         self.is(other)
     }
 
-    fn do_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
+    fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
         let result = vm._eq(self.clone(), other_key.clone())?;
         objbool::boolval(vm, result)
     }
@@ -307,8 +312,35 @@ impl DictKey for PyObjectRef {
 
 /// Implement trait for the str type, so that we can use strings
 /// to index dictionaries.
-impl DictKey for String {
-    fn do_hash(&self, _vm: &VirtualMachine) -> PyResult<HashValue> {
+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_string()).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)
+    }
+
+    fn do_is(self, _other: &PyObjectRef) -> bool {
+        // No matter who the other pyobject is, we are never the same thing, since
+        // we are a str, not a pyobject.
+        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.value == self)
+        } else {
+            // Fall back to PyString implementation.
+            let s = vm.new_str(self.to_string());
+            s.do_eq(vm, other_key)
+        }
+    }
+}
+
+impl DictKey for &String {
+    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_bigint().unwrap();
         let raw_hash = pyhash::hash_bigint(&raw_hash);
@@ -317,13 +349,13 @@ impl DictKey for String {
         Ok(hasher.finish() as HashValue)
     }
 
-    fn do_is(&self, _other: &PyObjectRef) -> bool {
+    fn do_is(self, _other: &PyObjectRef) -> bool {
         // No matter who the other pyobject is, we are never the same thing, since
         // we are a str, not a pyobject.
         false
     }
 
-    fn do_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
+    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.value == self)
         } else {
@@ -364,9 +396,9 @@ mod tests {
         assert_eq!(2, dict.len());
 
         assert_eq!(true, dict.contains(&vm, &key1).unwrap());
-        assert_eq!(true, dict.contains(&vm, &"x".to_string()).unwrap());
+        assert_eq!(true, dict.contains(&vm, "x").unwrap());
 
-        let val = dict.get(&vm, &"x".to_string()).unwrap().unwrap();
+        let val = dict.get(&vm, "x").unwrap().unwrap();
         vm._eq(val, value2)
             .expect("retrieved value must be equal to inserted value.");
     }
@@ -389,8 +421,8 @@ mod tests {
 
     fn check_hash_equivalence(text: &str) {
         let vm: VirtualMachine = Default::default();
-        let value1 = text.to_string();
-        let value2 = vm.new_str(value1.clone());
+        let value1 = text;
+        let value2 = vm.new_str(value1.to_string());
 
         let hash1 = value1.do_hash(&vm).expect("Hash should not fail.");
         let hash2 = value2.do_hash(&vm).expect("Hash should not fail.");

vm/src/frame.rs

@@ -285,12 +285,12 @@ impl Frame {
                         let dict: PyDictRef =
                             obj.downcast().expect("Need a dictionary to build a map.");
                         for (key, value) in dict {
-                            map_obj.set_item(key, value, vm).unwrap();
+                            map_obj.set_item(&key, value, vm).unwrap();
                         }
                     }
                 } else {
                     for (key, value) in self.pop_multiple(2 * size).into_iter().tuples() {
-                        map_obj.set_item(key, value, vm).unwrap();
+                        map_obj.set_item(&key, value, vm).unwrap();
                     }
                 }
 
@@ -946,14 +946,14 @@ impl Frame {
         let idx = self.pop_value();
         let obj = self.pop_value();
         let value = self.pop_value();
-        obj.set_item(idx, value, vm)?;
+        obj.set_item(&idx, value, vm)?;
         Ok(None)
     }
 
     fn execute_delete_subscript(&self, vm: &VirtualMachine) -> FrameResult {
         let idx = self.pop_value();
         let obj = self.pop_value();
-        obj.del_item(idx, vm)?;
+        obj.del_item(&idx, vm)?;
         Ok(None)
     }
 
@@ -1017,7 +1017,7 @@ impl Frame {
         let module = self
             .scope
             .globals
-            .get_item_option(&"__name__".to_string(), vm)?
+            .get_item_option("__name__", vm)?
             .unwrap_or_else(|| vm.get_none());
         vm.set_attr(&func_obj, "__module__", module)?;
         vm.set_attr(&func_obj, "__annotations__", annotations)?;
@@ -1062,7 +1062,7 @@ impl Frame {
                 bytecode::BinaryOperator::Divide => vm._truediv(a_ref, b_ref),
                 bytecode::BinaryOperator::FloorDivide => vm._floordiv(a_ref, b_ref),
                 // TODO: Subscript should probably have its own op
-                bytecode::BinaryOperator::Subscript => a_ref.get_item(b_ref, vm),
+                bytecode::BinaryOperator::Subscript => a_ref.get_item(&b_ref, vm),
                 bytecode::BinaryOperator::Modulo => vm._mod(a_ref, b_ref),
                 bytecode::BinaryOperator::Lshift => vm._lshift(a_ref, b_ref),
                 bytecode::BinaryOperator::Rshift => vm._rshift(a_ref, b_ref),

vm/src/obj/objdict.rs

@@ -200,7 +200,18 @@ impl PyDictRef {
         value: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
-        self.entries.borrow_mut().insert(vm, &key, value)
+        self.inner_setitem_fast(&key, value, vm)
+    }
+
+    /// Set item variant which can be called with multiple
+    /// key types, such as str to name a notable one.
+    fn inner_setitem_fast<K: DictKey + IntoPyObject + Copy>(
+        &self,
+        key: K,
+        value: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult<()> {
+        self.entries.borrow_mut().insert(vm, key, value)
     }
 
     #[cfg_attr(feature = "flame-it", flame("PyDictRef"))]
@@ -213,9 +224,9 @@ impl PyDictRef {
     }
 
     /// Return an optional inner item, or an error (can be key error as well)
-    fn inner_getitem_option<K: DictKey + IntoPyObject + Clone>(
+    fn inner_getitem_option<K: DictKey + IntoPyObject + Copy>(
         &self,
-        key: &K,
+        key: K,
         vm: &VirtualMachine,
     ) -> PyResult<Option<PyObjectRef>> {
         if let Some(value) = self.entries.borrow().get(vm, key)? {
@@ -224,9 +235,7 @@ impl PyDictRef {
 
         if let Some(method_or_err) = vm.get_method(self.clone().into_object(), "__missing__") {
             let method = method_or_err?;
-            Ok(Some(
-                vm.invoke(&method, vec![key.clone().into_pyobject(vm)?])?,
-            ))
+            Ok(Some(vm.invoke(&method, vec![key.into_pyobject(vm)?])?))
         } else {
             Ok(None)
         }
@@ -340,9 +349,9 @@ impl PyDictRef {
     /// python value, or None.
     /// Note that we can pass any type which implements the DictKey
     /// trait. Notable examples are String and PyObjectRef.
-    pub fn get_item_option<T: IntoPyObject + DictKey + Clone>(
+    pub fn get_item_option<T: IntoPyObject + DictKey + Copy>(
         &self,
-        key: &T,
+        key: T,
         vm: &VirtualMachine,
     ) -> PyResult<Option<PyObjectRef>> {
         // Test if this object is a true dict, or mabye a subclass?
@@ -365,7 +374,7 @@ impl PyDictRef {
         } else {
             // Fall back to full get_item with KeyError checking
 
-            match self.get_item(key.clone().into_pyobject(vm)?, vm) {
+            match self.get_item(key, vm) {
                 Ok(value) => Ok(Some(value)),
                 Err(exc) => {
                     if objtype::isinstance(&exc, &vm.ctx.exceptions.key_error) {
@@ -380,20 +389,26 @@ impl PyDictRef {
 }
 
 impl ItemProtocol for PyDictRef {
-    fn get_item<T: IntoPyObject>(&self, key: T, vm: &VirtualMachine) -> PyResult {
+    fn get_item<T: IntoPyObject + DictKey + Copy>(&self, key: T, vm: &VirtualMachine) -> PyResult {
         self.as_object().get_item(key, vm)
     }
 
-    fn set_item<T: IntoPyObject>(
+    fn set_item<T: IntoPyObject + DictKey + Copy>(
         &self,
         key: T,
         value: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult {
-        self.as_object().set_item(key, value, vm)
+        if self.typ().is(&vm.ctx.dict_type()) {
+            self.inner_setitem_fast(key, value, vm)
+                .map(|_| vm.ctx.none())
+        } else {
+            // Fall back to slow path if we are in a dict subclass:
+            self.as_object().set_item(key, value, vm)
+        }
     }
 
-    fn del_item<T: IntoPyObject>(&self, key: T, vm: &VirtualMachine) -> PyResult {
+    fn del_item<T: IntoPyObject + DictKey + Copy>(&self, key: T, vm: &VirtualMachine) -> PyResult {
         self.as_object().del_item(key, vm)
     }
 }

vm/src/obj/objobject.rs

@@ -78,7 +78,7 @@ fn object_setattr(
     }
 
     if let Some(ref dict) = obj.clone().dict {
-        dict.set_item(attr_name, value, vm)?;
+        dict.set_item(&attr_name.value, value, vm)?;
         Ok(())
     } else {
         Err(vm.new_attribute_error(format!(
@@ -99,7 +99,7 @@ fn object_delattr(obj: PyObjectRef, attr_name: PyStringRef, vm: &VirtualMachine)
     }
 
     if let Some(ref dict) = obj.dict {
-        dict.del_item(attr_name, vm)?;
+        dict.del_item(&attr_name.value, vm)?;
         Ok(())
     } else {
         Err(vm.new_attribute_error(format!(

vm/src/obj/objstr.rs

@@ -996,7 +996,7 @@ impl PyString {
 
         let mut translated = String::new();
         for c in self.value.chars() {
-            match table.get_item(c as u32, vm) {
+            match table.get_item(&(c as u32).into_pyobject(vm)?, vm) {
                 Ok(value) => {
                     if let Some(text) = value.payload::<PyString>() {
                         translated.push_str(&text.value);
@@ -1036,11 +1036,11 @@ impl PyString {
                 Ok(from_str) => {
                     if to_str.len(vm) == from_str.len(vm) {
                         for (c1, c2) in from_str.value.chars().zip(to_str.value.chars()) {
-                            new_dict.set_item(c1 as u32, vm.new_int(c2 as u32), vm)?;
+                            new_dict.set_item(&vm.new_int(c1 as u32), vm.new_int(c2 as u32), vm)?;
                         }
                         if let OptionalArg::Present(none_str) = none_str {
                             for c in none_str.value.chars() {
-                                new_dict.set_item(c as u32, vm.get_none(), vm)?;
+                                new_dict.set_item(&vm.new_int(c as u32), vm.get_none(), vm)?;
                             }
                         }
                         new_dict.into_pyobject(vm)
@@ -1061,11 +1061,15 @@ impl PyString {
                 Ok(dict) => {
                     for (key, val) in dict {
                         if let Some(num) = key.payload::<PyInt>() {
-                            new_dict.set_item(num.as_bigint().to_i32(), val, vm)?;
+                            new_dict.set_item(
+                                &num.as_bigint().to_i32().into_pyobject(vm)?,
+                                val,
+                                vm,
+                            )?;
                         } else if let Some(string) = key.payload::<PyString>() {
                             if string.len(vm) == 1 {
                                 let num_value = string.value.chars().next().unwrap() as u32;
-                                new_dict.set_item(num_value, val, vm)?;
+                                new_dict.set_item(&num_value.into_pyobject(vm)?, val, vm)?;
                             } else {
                                 return Err(vm.new_value_error(
                                     "string keys in translate table must be of length 1".to_owned(),

vm/src/py_serde.rs

@@ -211,7 +211,7 @@ impl<'de> Visitor<'de> for PyObjectDeserializer<'de> {
         // Although JSON keys must be strings, implementation accepts any keys
         // and can be reused by other deserializers without such limit
         while let Some((key_obj, value)) = access.next_entry_seed(self.clone(), self.clone())? {
-            dict.set_item(key_obj, value, self.vm).unwrap();
+            dict.set_item(&key_obj, value, self.vm).unwrap();
         }
         Ok(dict.into_object())
     }

vm/src/pyobject.rs

@@ -11,6 +11,7 @@ use num_complex::Complex64;
 use num_traits::{One, Zero};
 
 use crate::bytecode;
+use crate::dictdatatype::DictKey;
 use crate::exceptions;
 use crate::function::{IntoPyNativeFunc, PyFuncArgs};
 use crate::obj::objbuiltinfunc::PyBuiltinFunction;
@@ -764,15 +765,17 @@ impl<T> TypeProtocol for PyRef<T> {
     }
 }
 
+/// The python item protocol. Mostly applies to dictionaries.
+/// Allows getting, setting and deletion of keys-value pairs.
 pub trait ItemProtocol {
-    fn get_item<T: IntoPyObject>(&self, key: T, vm: &VirtualMachine) -> PyResult;
-    fn set_item<T: IntoPyObject>(
+    fn get_item<T: IntoPyObject + DictKey + Copy>(&self, key: T, vm: &VirtualMachine) -> PyResult;
+    fn set_item<T: IntoPyObject + DictKey + Copy>(
         &self,
         key: T,
         value: PyObjectRef,
         vm: &VirtualMachine,
     ) -> PyResult;
-    fn del_item<T: IntoPyObject>(&self, key: T, vm: &VirtualMachine) -> PyResult;
+    fn del_item<T: IntoPyObject + DictKey + Copy>(&self, key: T, vm: &VirtualMachine) -> PyResult;
 }
 
 impl ItemProtocol for PyObjectRef {
@@ -965,6 +968,12 @@ impl IntoPyObject for PyObjectRef {
     }
 }
 
+impl IntoPyObject for &PyObjectRef {
+    fn into_pyobject(self, _vm: &VirtualMachine) -> PyResult {
+        Ok(self.clone())
+    }
+}
+
 impl<T> IntoPyObject for PyResult<T>
 where
     T: IntoPyObject,

vm/src/scope.rs

@@ -133,12 +133,12 @@ impl NameProtocol for Scope {
     #[cfg_attr(feature = "flame-it", flame("Scope"))]
     fn load_name(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef> {
         for dict in self.locals.iter() {
-            if let Some(value) = dict.get_item_option(&name.to_string(), vm).unwrap() {
+            if let Some(value) = dict.get_item_option(name, vm).unwrap() {
                 return Some(value);
             }
         }
 
-        if let Some(value) = self.globals.get_item_option(&name.to_string(), vm).unwrap() {
+        if let Some(value) = self.globals.get_item_option(name, vm).unwrap() {
             return Some(value);
         }
 
@@ -148,7 +148,7 @@ impl NameProtocol for Scope {
     #[cfg_attr(feature = "flame-it", flame("Scope"))]
     fn load_cell(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef> {
         for dict in self.locals.iter().skip(1) {
-            if let Some(value) = dict.get_item_option(&name.to_string(), vm).unwrap() {
+            if let Some(value) = dict.get_item_option(name, vm).unwrap() {
                 return Some(value);
             }
         }
@@ -174,7 +174,7 @@ impl NameProtocol for Scope {
 
     #[cfg_attr(feature = "flame-it", flame("Scope"))]
     fn load_global(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef> {
-        self.globals.get_item_option(&name.to_string(), vm).unwrap()
+        self.globals.get_item_option(name, vm).unwrap()
     }
 
     fn store_global(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef) {

vm/src/stdlib/errno.rs

@@ -9,7 +9,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     for (name, code) in ERROR_CODES {
         let name = vm.new_str((*name).to_owned());
         let code = vm.ctx.new_int(*code);
-        errorcode.set_item(code.clone(), name.clone(), vm).unwrap();
+        errorcode.set_item(&code, name.clone(), vm).unwrap();
         vm.set_attr(&module, name, code).unwrap();
     }
     module

vm/src/vm.rs

@@ -452,14 +452,13 @@ impl VirtualMachine {
             || level != 0
             || objbool::boolval(self, from_list.clone()).unwrap_or(true);
 
-        let module = self.new_str(module.to_owned());
-
         let cached_module = if weird {
             None
         } else {
             let sys_modules = self.get_attribute(self.sys_module.clone(), "modules")?;
-            sys_modules.get_item(module.clone(), self).ok()
+            sys_modules.get_item(module, self).ok()
         };
+
         match cached_module {
             Some(module) => Ok(module),
             None => {
@@ -478,7 +477,7 @@ impl VirtualMachine {
                 self.invoke(
                     &import_func,
                     vec![
-                        module,
+                        self.new_str(module.to_owned()),
                         globals,
                         locals,
                         from_list.clone(),