[pypy-svn] pypy extend-rweakdict: RWeakValueDictionary now accept different kinds of keys.

amauryfa commits-noreply at bitbucket.org
Sun Mar 13 23:41:35 CET 2011


Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: extend-rweakdict
Changeset: r42578:76f8aa04732d
Date: 2011-03-13 22:18 +0100
http://bitbucket.org/pypy/pypy/changeset/76f8aa04732d/

Log:	RWeakValueDictionary now accept different kinds of keys. Only tested
	with str and int.

diff --git a/pypy/rlib/rweakref.py b/pypy/rlib/rweakref.py
--- a/pypy/rlib/rweakref.py
+++ b/pypy/rlib/rweakref.py
@@ -8,9 +8,7 @@
 
 
 class RWeakValueDictionary(object):
-    """A limited dictionary containing weak values.
-    Only supports string keys.
-    """
+    """A dictionary containing weak values."""
 
     def __init__(self, keyclass, valueclass):
         self._dict = weakref.WeakValueDictionary()
@@ -70,18 +68,19 @@
 class SomeWeakValueDict(annmodel.SomeObject):
     knowntype = RWeakValueDictionary
 
-    def __init__(self, valueclassdef):
+    def __init__(self, s_key, valueclassdef):
+        self.s_key = s_key
         self.valueclassdef = valueclassdef
 
     def rtyper_makerepr(self, rtyper):
         from pypy.rlib import _rweakvaldict
-        return _rweakvaldict.WeakValueDictRepr(rtyper)
+        return _rweakvaldict.WeakValueDictRepr(rtyper,
+                                               rtyper.makerepr(self.s_key))
 
     def rtyper_makekey_ex(self, rtyper):
         return self.__class__,
 
     def method_get(self, s_key):
-        assert annmodel.SomeString(can_be_None=True).contains(s_key)
         return annmodel.SomeInstance(self.valueclassdef, can_be_None=True)
 
     def method_set(self, s_key, s_value):
@@ -91,18 +90,23 @@
 class __extend__(pairtype(SomeWeakValueDict, SomeWeakValueDict)):
     def union((s_wvd1, s_wvd2)):
         if s_wvd1.valueclassdef is not s_wvd2.valueclassdef:
-            return SomeObject() # not the same class! complain...
-        return SomeWeakValueDict(s_wvd1.valueclassdef)
+            return annmodel.SomeObject() # not the same class! complain...
+        s_key = annmodel.unionof(s_wvd1.s_key, s_wvd2.s_key)
+        return SomeWeakValueDict(s_key, s_wvd1.valueclassdef)
 
 class Entry(extregistry.ExtRegistryEntry):
     _about_ = RWeakValueDictionary
 
     def compute_result_annotation(self, s_keyclass, s_valueclass):
-        return SomeWeakValueDict(_getclassdef(s_valueclass))
+        assert s_keyclass.is_constant()
+        s_key = self.bookkeeper.immutablevalue(s_keyclass.const())
+        return SomeWeakValueDict(
+            s_key,
+            _getclassdef(s_valueclass))
 
     def specialize_call(self, hop):
         from pypy.rlib import _rweakvaldict
-        return _rweakvaldict.specialize_make_weakdict(hop)
+        return _rweakvaldict.specialize_make_weakdict(hop, hop.r_result.traits)
 
 class Entry(extregistry.ExtRegistryEntry):
     _type_ = RWeakValueDictionary
@@ -110,7 +114,9 @@
     def compute_annotation(self):
         bk = self.bookkeeper
         x = self.instance
-        return SomeWeakValueDict(bk.getuniqueclassdef(x._valueclass))
+        return SomeWeakValueDict(
+            bk.immutablevalue(x._keyclass()),
+            bk.getuniqueclassdef(x._valueclass))
 
 def _getclassdef(s_instance):
     assert isinstance(s_instance, annmodel.SomePBC)

diff --git a/pypy/rlib/_rweakvaldict.py b/pypy/rlib/_rweakvaldict.py
--- a/pypy/rlib/_rweakvaldict.py
+++ b/pypy/rlib/_rweakvaldict.py
@@ -3,15 +3,18 @@
 from pypy.rpython.lltypesystem.llmemory import weakref_create, weakref_deref
 from pypy.rpython.lltypesystem.lloperation import llop
 from pypy.rpython.rclass import getinstancerepr
+from pypy.rpython.rint import signed_repr
 from pypy.rpython.rmodel import Repr
 from pypy.rlib.rweakref import RWeakValueDictionary
 from pypy.rlib import jit
 
 
 class WeakValueDictRepr(Repr):
-    def __init__(self, rtyper):
+    def __init__(self, rtyper, r_key):
         self.rtyper = rtyper
-        self.lowleveltype = lltype.Ptr(WEAKDICT)
+        self.r_key = r_key
+        self.traits = make_WEAKDICT(r_key)
+        self.lowleveltype = lltype.Ptr(self.traits.WEAKDICT)
         self.dict_cache = {}
 
     def convert_const(self, weakdict):
@@ -23,150 +26,163 @@
             return self.dict_cache[key]
         except KeyError:
             self.setup()
-            l_dict = ll_new_weakdict()
+            l_dict = self.traits.ll_new_weakdict()
             self.dict_cache[key] = l_dict
             bk = self.rtyper.annotator.bookkeeper
             classdef = bk.getuniqueclassdef(weakdict._valueclass)
-            r_key = rstr.string_repr
             r_value = getinstancerepr(self.rtyper, classdef)
             for dictkey, dictvalue in weakdict._dict.items():
-                llkey = r_key.convert_const(dictkey)
+                llkey = self.r_key.convert_const(dictkey)
                 llvalue = r_value.convert_const(dictvalue)
                 if llvalue:
                     llvalue = lltype.cast_pointer(rclass.OBJECTPTR, llvalue)
-                    ll_set_nonnull(l_dict, llkey, llvalue)
+                    self.traits.ll_set_nonnull(l_dict, llkey, llvalue)
             return l_dict
 
     def rtype_method_get(self, hop):
-        v_d, v_key = hop.inputargs(self, rstr.string_repr)
+        
+        v_d, v_key = hop.inputargs(self, self.r_key)
         hop.exception_cannot_occur()
-        v_result = hop.gendirectcall(ll_get, v_d, v_key)
+        v_result = hop.gendirectcall(self.traits.ll_get, v_d, v_key)
         v_result = hop.genop("cast_pointer", [v_result],
                              resulttype=hop.r_result.lowleveltype)
         return v_result
 
     def rtype_method_set(self, hop):
         r_object = getinstancerepr(self.rtyper, None)
-        v_d, v_key, v_value = hop.inputargs(self, rstr.string_repr,
-                                            r_object)
+        v_d, v_key, v_value = hop.inputargs(self, self.r_key, r_object)
         hop.exception_cannot_occur()
         if hop.args_s[2].is_constant() and hop.args_s[2].const is None:
-            hop.gendirectcall(ll_set_null, v_d, v_key)
+            hop.gendirectcall(self.traits.ll_set_null, v_d, v_key)
         else:
-            hop.gendirectcall(ll_set, v_d, v_key, v_value)
+            hop.gendirectcall(self.traits.ll_set, v_d, v_key, v_value)
 
 
-def specialize_make_weakdict(hop):
+def specialize_make_weakdict(hop, traits):
     hop.exception_cannot_occur()
-    v_d = hop.gendirectcall(ll_new_weakdict)
+    v_d = hop.gendirectcall(traits.ll_new_weakdict)
     return v_d
 
 # ____________________________________________________________
 
+def make_WEAKDICT(r_key):
+    KEY = r_key.lowleveltype
+    ll_keyhash = r_key.get_ll_hash_function()
+    if isinstance(KEY, lltype.Ptr):
+        zero_key = r_key.convert_const(None)
+    else:
+        zero_key = r_key.convert_const(0)
 
-WEAKDICTENTRY = lltype.Struct("weakdictentry",
-                              ("key", lltype.Ptr(rstr.STR)),
-                              ("value", llmemory.WeakRefPtr))
+    WEAKDICTENTRY = lltype.Struct("weakdictentry",
+                                  ("key", KEY),
+                                  ("value", llmemory.WeakRefPtr))
 
-def ll_valid(entries, i):
-    value = entries[i].value
-    return bool(value) and bool(weakref_deref(rclass.OBJECTPTR, value))
+    def ll_valid(entries, i):
+        value = entries[i].value
+        return bool(value) and bool(weakref_deref(rclass.OBJECTPTR, value))
 
-def ll_everused(entries, i):
-    return bool(entries[i].value)
+    def ll_everused(entries, i):
+        return bool(entries[i].value)
 
-def ll_hash(entries, i):
-    return str_fasthashfn(entries[i].key)
-str_fasthashfn = rstr.string_repr.get_ll_fasthash_function()
+    def ll_hash(entries, i):
+        return fasthashfn(entries[i].key)
+    fasthashfn = r_key.get_ll_fasthash_function()
 
-entrymeths = {
-    'allocate': lltype.typeMethod(rdict._ll_malloc_entries),
-    'delete': rdict._ll_free_entries,
-    'valid': ll_valid,
-    'everused': ll_everused,
-    'hash': ll_hash,
-    }
-WEAKDICTENTRYARRAY = lltype.GcArray(WEAKDICTENTRY,
-                                    adtmeths=entrymeths,
-                                    hints={'weakarray': 'value'})
-# NB. the 'hints' is not used so far ^^^
+    entrymeths = {
+        'allocate': lltype.typeMethod(rdict._ll_malloc_entries),
+        'delete': rdict._ll_free_entries,
+        'valid': ll_valid,
+        'everused': ll_everused,
+        'hash': ll_hash,
+        }
+    WEAKDICTENTRYARRAY = lltype.GcArray(WEAKDICTENTRY,
+                                        adtmeths=entrymeths,
+                                        hints={'weakarray': 'value'})
+    # NB. the 'hints' is not used so far ^^^
 
-ll_strhash = rstr.LLHelpers.ll_strhash
+    class Traits:
+        @staticmethod
+        @jit.dont_look_inside
+        def ll_new_weakdict():
+            d = lltype.malloc(Traits.WEAKDICT)
+            d.entries = Traits.WEAKDICT.entries.TO.allocate(rdict.DICT_INITSIZE)
+            d.num_items = 0
+            d.num_pristine_entries = rdict.DICT_INITSIZE
+            return d
 
- at jit.dont_look_inside
-def ll_new_weakdict():
-    d = lltype.malloc(WEAKDICT)
-    d.entries = WEAKDICT.entries.TO.allocate(rdict.DICT_INITSIZE)
-    d.num_items = 0
-    d.num_pristine_entries = rdict.DICT_INITSIZE
-    return d
+        @staticmethod
+        @jit.dont_look_inside
+        def ll_get(d, llkey):
+            hash = ll_keyhash(llkey)
+            i = rdict.ll_dict_lookup(d, llkey, hash)
+            #llop.debug_print(lltype.Void, i, 'get')
+            valueref = d.entries[i].value
+            if valueref:
+                return weakref_deref(rclass.OBJECTPTR, valueref)
+            else:
+                return lltype.nullptr(rclass.OBJECTPTR.TO)
 
- at jit.dont_look_inside
-def ll_get(d, llkey):
-    hash = ll_strhash(llkey)
-    i = rdict.ll_dict_lookup(d, llkey, hash)
-    #llop.debug_print(lltype.Void, i, 'get')
-    valueref = d.entries[i].value
-    if valueref:
-        return weakref_deref(rclass.OBJECTPTR, valueref)
-    else:
-        return lltype.nullptr(rclass.OBJECTPTR.TO)
+        @staticmethod
+        @jit.dont_look_inside
+        def ll_set(d, llkey, llvalue):
+            if llvalue:
+                Traits.ll_set_nonnull(d, llkey, llvalue)
+            else:
+                Traits.ll_set_null(d, llkey)
 
- at jit.dont_look_inside
-def ll_set(d, llkey, llvalue):
-    if llvalue:
-        ll_set_nonnull(d, llkey, llvalue)
-    else:
-        ll_set_null(d, llkey)
+        @staticmethod
+        @jit.dont_look_inside
+        def ll_set_nonnull(d, llkey, llvalue):
+            hash = ll_keyhash(llkey)
+            valueref = weakref_create(llvalue)    # GC effects here, before the rest
+            i = rdict.ll_dict_lookup(d, llkey, hash)
+            everused = d.entries.everused(i)
+            d.entries[i].key = llkey
+            d.entries[i].value = valueref
+            #llop.debug_print(lltype.Void, i, 'stored')
+            if not everused:
+                d.num_pristine_entries -= 1
+                if d.num_pristine_entries * 3 <= len(d.entries):
+                    #llop.debug_print(lltype.Void, 'RESIZE')
+                    Traits.ll_weakdict_resize(d)
 
- at jit.dont_look_inside
-def ll_set_nonnull(d, llkey, llvalue):
-    hash = ll_strhash(llkey)
-    valueref = weakref_create(llvalue)    # GC effects here, before the rest
-    i = rdict.ll_dict_lookup(d, llkey, hash)
-    everused = d.entries.everused(i)
-    d.entries[i].key = llkey
-    d.entries[i].value = valueref
-    #llop.debug_print(lltype.Void, i, 'stored')
-    if not everused:
-        d.num_pristine_entries -= 1
-        if d.num_pristine_entries * 3 <= len(d.entries):
-            #llop.debug_print(lltype.Void, 'RESIZE')
-            ll_weakdict_resize(d)
+        @staticmethod
+        @jit.dont_look_inside
+        def ll_set_null(d, llkey):
+            hash = ll_keyhash(llkey)
+            i = rdict.ll_dict_lookup(d, llkey, hash)
+            if d.entries.everused(i):
+                # If the entry was ever used, clean up its key and value.
+                # We don't store a NULL value, but a dead weakref, because
+                # the entry must still be marked as everused().
+                d.entries[i].value = llmemory.dead_wref
+                d.entries[i].key = zero_key
+                #llop.debug_print(lltype.Void, i, 'zero')
 
- at jit.dont_look_inside
-def ll_set_null(d, llkey):
-    hash = ll_strhash(llkey)
-    i = rdict.ll_dict_lookup(d, llkey, hash)
-    if d.entries.everused(i):
-        # If the entry was ever used, clean up its key and value.
-        # We don't store a NULL value, but a dead weakref, because
-        # the entry must still be marked as everused().
-        d.entries[i].value = llmemory.dead_wref
-        d.entries[i].key = lltype.nullptr(rstr.STR)
-        #llop.debug_print(lltype.Void, i, 'zero')
+        @staticmethod
+        def ll_weakdict_resize(d):
+            # first set num_items to its correct, up-to-date value
+            entries = d.entries
+            num_items = 0
+            for i in range(len(entries)):
+                if entries.valid(i):
+                    num_items += 1
+            d.num_items = num_items
+            rdict.ll_dict_resize(d)
 
-def ll_weakdict_resize(d):
-    # first set num_items to its correct, up-to-date value
-    entries = d.entries
-    num_items = 0
-    for i in range(len(entries)):
-        if entries.valid(i):
-            num_items += 1
-    d.num_items = num_items
-    rdict.ll_dict_resize(d)
+        ll_keyeq = lltype.staticAdtMethod(r_key.get_ll_eq_function())
 
-str_keyeq = lltype.staticAdtMethod(rstr.string_repr.get_ll_eq_function())
+        dictmeths = {
+            'll_get': ll_get,
+            'll_set': ll_set,
+            'keyeq': ll_keyeq,
+            'paranoia': False,
+            }
 
-dictmeths = {
-    'll_get': ll_get,
-    'll_set': ll_set,
-    'keyeq': str_keyeq,
-    'paranoia': False,
-    }
+        WEAKDICT = lltype.GcStruct("weakvaldict",
+                                   ("num_items", lltype.Signed),
+                                   ("num_pristine_entries", lltype.Signed),
+                                   ("entries", lltype.Ptr(WEAKDICTENTRYARRAY)),
+                                   adtmeths=dictmeths)
 
-WEAKDICT = lltype.GcStruct("weakvaldict",
-                           ("num_items", lltype.Signed),
-                           ("num_pristine_entries", lltype.Signed),
-                           ("entries", lltype.Ptr(WEAKDICTENTRYARRAY)),
-                           adtmeths=dictmeths)
+    return Traits

diff --git a/pypy/rlib/test/test_rweakvaldict.py b/pypy/rlib/test/test_rweakvaldict.py
--- a/pypy/rlib/test/test_rweakvaldict.py
+++ b/pypy/rlib/test/test_rweakvaldict.py
@@ -10,65 +10,78 @@
     pass
 
 
-def make_test(loop=100):
+def make_test(loop=100, keyclass=str):
+    if keyclass is str:
+        make_key = str
+        keys = ["abc", "def", "ghi", "hello"]
+    elif keyclass is int:
+        make_key = int
+        keys = [123, 456, 789, 1234]
+
     def g(d):
-        assert d.get("hello") is None
+        assert d.get(keys[3]) is None
         x1 = X(); x2 = X(); x3 = X()
-        d.set("abc", x1)
-        d.set("def", x2)
-        d.set("ghi", x3)
-        assert d.get("abc") is x1
-        assert d.get("def") is x2
-        assert d.get("ghi") is x3
-        assert d.get("hello") is None
+        d.set(keys[0], x1)
+        d.set(keys[1], x2)
+        d.set(keys[2], x3)
+        assert d.get(keys[0]) is x1
+        assert d.get(keys[1]) is x2
+        assert d.get(keys[2]) is x3
+        assert d.get(keys[3]) is None
         return x1, x3    # x2 dies
     def f():
-        d = RWeakValueDictionary(str, X)
+        d = RWeakValueDictionary(keyclass, X)
         x1, x3 = g(d)
         rgc.collect(); rgc.collect()
-        assert d.get("abc") is x1
-        assert d.get("def") is None
-        assert d.get("ghi") is x3
-        assert d.get("hello") is None
-        d.set("abc", None)
-        assert d.get("abc") is None
-        assert d.get("def") is None
-        assert d.get("ghi") is x3
-        assert d.get("hello") is None
+        assert d.get(keys[0]) is x1
+        assert d.get(keys[1]) is None
+        assert d.get(keys[2]) is x3
+        assert d.get(keys[3]) is None
+        d.set(keys[0], None)
+        assert d.get(keys[0]) is None
+        assert d.get(keys[1]) is None
+        assert d.get(keys[2]) is x3
+        assert d.get(keys[3]) is None
         # resizing should also work
         for i in range(loop):
-            d.set(str(i), x1)
+            d.set(make_key(i), x1)
         for i in range(loop):
-            assert d.get(str(i)) is x1
-        assert d.get("abc") is None
-        assert d.get("def") is None
-        assert d.get("ghi") is x3
-        assert d.get("hello") is None
+            assert d.get(make_key(i)) is x1
+        assert d.get(keys[0]) is None
+        assert d.get(keys[1]) is None
+        assert d.get(keys[2]) is x3
+        assert d.get(keys[3]) is None
         # a subclass
         y = Y()
-        d.set("hello", y)
-        assert d.get("hello") is y
+        d.set(keys[3], y)
+        assert d.get(keys[3]) is y
         # storing a lot of Nones
         for i in range(loop, loop*2-5):
-            d.set('%dfoobar' % i, x1)
+            d.set(make_key(1000 + i), x1)
         for i in range(loop):
-            d.set(str(i), None)
+            d.set(make_key(i), None)
         for i in range(loop):
-            assert d.get(str(i)) is None
-        assert d.get("abc") is None
-        assert d.get("def") is None
-        assert d.get("ghi") is x3
-        assert d.get("hello") is y
+            assert d.get(make_key(i)) is None
+        assert d.get(keys[0]) is None
+        assert d.get(keys[1]) is None
+        assert d.get(keys[2]) is x3
+        assert d.get(keys[3]) is y
         for i in range(loop, loop*2-5):
-            assert d.get('%dfoobar' % i) is x1
+            assert d.get(make_key(1000 + i)) is x1
     return f
 
 def test_RWeakValueDictionary():
     make_test()()
 
+def test_RWeakValueDictionary_int():
+    make_test(keyclass=int)()
+
 def test_rpython_RWeakValueDictionary():
     interpret(make_test(loop=12), [])
 
+def test_rpython_RWeakValueDictionary_int():
+    interpret(make_test(loop=12, keyclass=int), [])
+
 def test_rpython_prebuilt():
     d = RWeakValueDictionary(str, X)
     living = [X() for i in range(8)]


More information about the Pypy-commit mailing list