diff --git a/src/main.rs b/src/main.rs index 899188d04..7e62ff914 100644 --- a/src/main.rs +++ b/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(); } } diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs index 4442ea7c6..5f086c2dc 100644 --- a/vm/src/dictdatatype.rs +++ b/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 Dict { &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 Dict { } /// Store a key - pub fn insert(&mut self, vm: &VirtualMachine, key: &PyObjectRef, value: T) -> PyResult<()> { + pub fn insert( + &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 Dict { 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(&self, vm: &VirtualMachine, key: &K) -> PyResult { + pub fn contains(&self, vm: &VirtualMachine, key: K) -> PyResult { if let LookupResult::Existing(_) = self.lookup(vm, key)? { Ok(true) } else { @@ -113,7 +118,7 @@ impl Dict { /// Retrieve a key #[cfg_attr(feature = "flame-it", flame("Dict"))] - pub fn get(&self, vm: &VirtualMachine, key: &K) -> PyResult> { + pub fn get(&self, vm: &VirtualMachine, key: K) -> PyResult> { if let LookupResult::Existing(index) = self.lookup(vm, key)? { Ok(Some(self.unchecked_get(index))) } else { @@ -156,7 +161,7 @@ impl Dict { 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 Dict { /// Lookup the index for the given key. #[cfg_attr(feature = "flame-it", flame("Dict"))] - fn lookup(&self, vm: &VirtualMachine, key: &K) -> PyResult { + fn lookup(&self, vm: &VirtualMachine, key: K) -> PyResult { let hash_value = key.do_hash(vm)?; let perturb = hash_value; let mut hash_index: HashIndex = hash_value; @@ -244,7 +249,7 @@ impl Dict { } /// Retrieve and delete a key - pub fn pop(&mut self, vm: &VirtualMachine, key: &K) -> PyResult> { + pub fn pop(&mut self, vm: &VirtualMachine, key: K) -> PyResult> { 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; - fn do_is(&self, other: &PyObjectRef) -> bool; - fn do_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult; + fn do_hash(self, vm: &VirtualMachine) -> PyResult; + fn do_is(self, other: &PyObjectRef) -> bool; + fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult; } /// 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 { +impl DictKey for &PyObjectRef { + fn do_hash(self, vm: &VirtualMachine) -> PyResult { 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 { + fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult { 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 { +impl DictKey for &str { + fn do_hash(self, _vm: &VirtualMachine) -> PyResult { + // 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 { + if let Some(py_str_value) = other_key.payload::() { + 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 { // 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 { + fn do_eq(self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult { if let Some(py_str_value) = other_key.payload::() { 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."); diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 418f80576..f205f4089 100644 --- a/vm/src/frame.rs +++ b/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), diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs index bb1f2d1b8..f645902d6 100644 --- a/vm/src/obj/objdict.rs +++ b/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( + &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( + fn inner_getitem_option( &self, - key: &K, + key: K, vm: &VirtualMachine, ) -> PyResult> { 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( + pub fn get_item_option( &self, - key: &T, + key: T, vm: &VirtualMachine, ) -> PyResult> { // 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(&self, key: T, vm: &VirtualMachine) -> PyResult { + fn get_item(&self, key: T, vm: &VirtualMachine) -> PyResult { self.as_object().get_item(key, vm) } - fn set_item( + fn set_item( &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(&self, key: T, vm: &VirtualMachine) -> PyResult { + fn del_item(&self, key: T, vm: &VirtualMachine) -> PyResult { self.as_object().del_item(key, vm) } } diff --git a/vm/src/obj/objobject.rs b/vm/src/obj/objobject.rs index 87b1da34a..70e9b5fd3 100644 --- a/vm/src/obj/objobject.rs +++ b/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!( diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs index 96949553c..15e1a9ac0 100644 --- a/vm/src/obj/objstr.rs +++ b/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::() { 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::() { - 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::() { 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(), diff --git a/vm/src/py_serde.rs b/vm/src/py_serde.rs index 5e43c43a7..92200da5b 100644 --- a/vm/src/py_serde.rs +++ b/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()) } diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs index 05049528c..85165466d 100644 --- a/vm/src/pyobject.rs +++ b/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 TypeProtocol for PyRef { } } +/// The python item protocol. Mostly applies to dictionaries. +/// Allows getting, setting and deletion of keys-value pairs. pub trait ItemProtocol { - fn get_item(&self, key: T, vm: &VirtualMachine) -> PyResult; - fn set_item( + fn get_item(&self, key: T, vm: &VirtualMachine) -> PyResult; + fn set_item( &self, key: T, value: PyObjectRef, vm: &VirtualMachine, ) -> PyResult; - fn del_item(&self, key: T, vm: &VirtualMachine) -> PyResult; + fn del_item(&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 IntoPyObject for PyResult where T: IntoPyObject, diff --git a/vm/src/scope.rs b/vm/src/scope.rs index e985343d5..9e5baadb7 100644 --- a/vm/src/scope.rs +++ b/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 { 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 { 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 { - 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) { diff --git a/vm/src/stdlib/errno.rs b/vm/src/stdlib/errno.rs index 6f82351f2..6967c3a4c 100644 --- a/vm/src/stdlib/errno.rs +++ b/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 diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 94797f1ff..ccb18f2d9 100644 --- a/vm/src/vm.rs +++ b/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(),