Implement do_eq for String impl of DictKey. Also move get_item_option to dict type.

windelbouwman · windelbouwman · commit f2a70bb08805 · 2019-07-28T20:39:22.000+02:00
diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs
@@ -1,4 +1,5 @@
 use crate::obj::objbool;
+use crate::obj::objstr::PyString;
 use crate::pyhash;
 use crate::pyobject::{IdProtocol, PyObjectRef, PyResult};
 use crate::vm::VirtualMachine;
@@ -308,6 +309,7 @@ impl DictKey for PyObjectRef {
 /// to index dictionaries.
 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);
         let mut hasher = DefaultHasher::new();
@@ -322,9 +324,13 @@ impl DictKey for String {
     }
 
     fn do_eq(&self, vm: &VirtualMachine, other_key: &PyObjectRef) -> PyResult<bool> {
-        // Fall back to PyString implementation.
-        let s = vm.new_str(self.to_string());
-        s.do_eq(vm, other_key)
+        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)
+        }
     }
 }
 
diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs
@@ -11,6 +11,7 @@ use crate::vm::{ReprGuard, VirtualMachine};
 use super::objbool;
 use super::objiter;
 use super::objstr;
+use super::objtype;
 use crate::dictdatatype;
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::PyClassImpl;
@@ -316,6 +317,23 @@ impl PyDictRef {
     pub fn size(&self) -> dictdatatype::DictSize {
         self.entries.borrow().size()
     }
+
+    pub fn get_item_option<T: IntoPyObject>(
+        &self,
+        key: T,
+        vm: &VirtualMachine,
+    ) -> PyResult<Option<PyObjectRef>> {
+        match self.get_item(key, vm) {
+            Ok(value) => Ok(Some(value)),
+            Err(exc) => {
+                if objtype::isinstance(&exc, &vm.ctx.exceptions.key_error) {
+                    Ok(None)
+                } else {
+                    Err(exc)
+                }
+            }
+        }
+    }
 }
 
 impl ItemProtocol for PyDictRef {
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -1044,24 +1044,6 @@ pub trait ItemProtocol {
         vm: &VirtualMachine,
     ) -> PyResult;
     fn del_item<T: IntoPyObject>(&self, key: T, vm: &VirtualMachine) -> PyResult;
-
-    #[cfg_attr(feature = "flame-it", flame("ItemProtocol"))]
-    fn get_item_option<T: IntoPyObject>(
-        &self,
-        key: T,
-        vm: &VirtualMachine,
-    ) -> PyResult<Option<PyObjectRef>> {
-        match self.get_item(key, vm) {
-            Ok(value) => Ok(Some(value)),
-            Err(exc) => {
-                if objtype::isinstance(&exc, &vm.ctx.exceptions.key_error) {
-                    Ok(None)
-                } else {
-                    Err(exc)
-                }
-            }
-        }
-    }
 }
 
 impl ItemProtocol for PyObjectRef {
