[pypy-commit] pypy keys_with_hash: in-progress

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: keys_with_hash
Changeset: r79331:d7605de37f95
Date: 2015-08-31 19:57 +0200
http://bitbucket.org/pypy/pypy/changeset/d7605de37f95/

Log:	in-progress

diff --git a/rpython/annotator/unaryop.py b/rpython/annotator/unaryop.py
--- a/rpython/annotator/unaryop.py
+++ b/rpython/annotator/unaryop.py
@@ -369,14 +369,19 @@
         return None    # r_dict: can throw anything
     return []          # else: no possible exception
 
- at op.contains.register(SomeDict)
-def contains_SomeDict(annotator, dct, element):
-    annotator.annotation(dct).dictdef.generalize_key(annotator.annotation(element))
-    if annotator.annotation(dct)._is_empty():
+# also used for objectmodel.contains_with_hash()
+def dict_contains(s_dct, s_element):
+    s_dct.dictdef.generalize_key(s_element)
+    if s_dct._is_empty():
         s_bool = SomeBool()
         s_bool.const = False
         return s_bool
     return s_Bool
+
+ at op.contains.register(SomeDict)
+def contains_SomeDict(annotator, dct, element):
+    return dict_contains(annotator.annotation(dct),
+                         annotator.annotation(element))
 contains_SomeDict.can_only_throw = _can_only_throw
 
 class __extend__(SomeDict):
diff --git a/rpython/rlib/objectmodel.py b/rpython/rlib/objectmodel.py
--- a/rpython/rlib/objectmodel.py
+++ b/rpython/rlib/objectmodel.py
@@ -792,7 +792,14 @@
     """Iterates (key, hash) pairs without recomputing the hash."""
     assert not we_are_translated()    # this code is only before translation
     for k in d:
-        yield k, compute_hash(k)
+        yield k, hash(k)
+
+def contains_with_hash(d, key, h):
+    """Same as 'key in d'.  The extra argument is the hash.  Use this only
+    if you got the hash just now from some other ..._with_hash() function."""
+    assert not we_are_translated()    # this code is only before translation
+    assert hash(key) == h
+    return key in d
 
 class Entry(ExtRegistryEntry):
     _about_ = iterkeys_with_hash
@@ -810,6 +817,27 @@
         hop.exception_cannot_occur()
         return DictIteratorRepr(hop.args_r[0], "keys_with_hash").newiter(hop)
 
+class Entry(ExtRegistryEntry):
+    _about_ = contains_with_hash
+
+    def compute_result_annotation(self, s_d, s_key, s_hash):
+        from rpython.annotator.model import s_Bool, SomeDict, s_None
+        from rpython.annotator.unaryop import dict_contains
+        if isinstance(s_d, SomeDict):
+            return dict_contains(s_d, s_key)
+        if s_None.contains(s_d):
+            return None
+        raise Exception("contains_with_hash(x, ...): x not a dict")
+
+    def specialize_call(self, hop):
+        from rpython.rtyper.lltypesystem import lltype
+        from rpython.rtyper.lltypesystem.rdict import ll_contains_with_hash
+        r_dict = hop.args_r[0]
+        v_dict, v_key, v_hash = hop.inputargs(r_dict, r_dict.key_repr,
+                                              lltype.Signed)
+        hop.exception_cannot_occur()
+        return hop.gendirectcall(ll_contains_with_hash, v_dict, v_key, v_hash)
+
 # ____________________________________________________________
 
 def import_from_mixin(M, special_methods=['__init__', '__del__']):
diff --git a/rpython/rlib/test/test_objectmodel.py b/rpython/rlib/test/test_objectmodel.py
--- a/rpython/rlib/test/test_objectmodel.py
+++ b/rpython/rlib/test/test_objectmodel.py
@@ -608,6 +608,16 @@
     r = interpret(f, [29])
     assert r == 0.0
 
+def test_contains_with_hash():
+    def f(i):
+        d = {i+.5: 5}
+        assert contains_with_hash(d, i+.5, compute_hash(i+.5))
+        assert not contains_with_hash(d, i+.3, compute_hash(i+.3))
+        return 0
+
+    f(29)
+    interpret(f, [29])
+
 def test_import_from_mixin():
     class M:    # old-style
         def f(self): pass
diff --git a/rpython/rtyper/lltypesystem/rdict.py b/rpython/rtyper/lltypesystem/rdict.py
--- a/rpython/rtyper/lltypesystem/rdict.py
+++ b/rpython/rtyper/lltypesystem/rdict.py
@@ -871,6 +871,11 @@
     i = ll_dict_lookup(d, key, d.keyhash(key))
     return not i & HIGHEST_BIT
 
+# for objectmodel.contains_with_hash()
+def ll_contains_with_hash(d, key, hash):
+    i = ll_dict_lookup(d, key, hash)
+    return not i & HIGHEST_BIT
+
 POPITEMINDEX = lltype.Struct('PopItemIndex', ('nextindex', lltype.Signed))
 global_popitem_index = lltype.malloc(POPITEMINDEX, zero=True, immortal=True)