Commits

Armin Rigo committed 711f53c

Refactor the implementation of the dict iterators in RPython. The goal
is to avoid the allocation of a tuple in iteritems().next().

  • Participants
  • Parent commits af210d9

Comments (0)

Files changed (5)

File rpython/jit/codewriter/support.py

     _ll_1_dict_values.need_result_type = True
     _ll_1_dict_items .need_result_type = True
 
-    _dictnext_keys   = staticmethod(ll_rdict.ll_dictnext_group['keys'])
-    _dictnext_values = staticmethod(ll_rdict.ll_dictnext_group['values'])
-    _dictnext_items  = staticmethod(ll_rdict.ll_dictnext_group['items'])
-
-    def _ll_1_dictiter_nextkeys(iter):
-        return LLtypeHelpers._dictnext_keys(None, iter)
-    def _ll_1_dictiter_nextvalues(iter):
-        return LLtypeHelpers._dictnext_values(None, iter)
-    def _ll_1_dictiter_nextitems(RES, iter):
-        return LLtypeHelpers._dictnext_items(lltype.Ptr(RES), iter)
-    _ll_1_dictiter_nextitems.need_result_type = True
-
+    _ll_1_dictiter_next = ll_rdict._ll_dictnext
     _ll_1_dict_resize = ll_rdict.ll_dict_resize
 
     # ---------- ordered dict ----------
     _ll_1_odict_values.need_result_type = True
     _ll_1_odict_items .need_result_type = True
 
-    _odictnext_keys   = staticmethod(rordereddict.ll_dictnext_group['keys'])
-    _odictnext_values = staticmethod(rordereddict.ll_dictnext_group['values'])
-    _odictnext_items  = staticmethod(rordereddict.ll_dictnext_group['items'])
-
-    def _ll_1_odictiter_nextkeys(iter):
-        return LLtypeHelpers._odictnext_keys(None, iter)
-    def _ll_1_odictiter_nextvalues(iter):
-        return LLtypeHelpers._odictnext_values(None, iter)
-    def _ll_1_odictiter_nextitems(RES, iter):
-        return LLtypeHelpers._odictnext_items(lltype.Ptr(RES), iter)
-    _ll_1_odictiter_nextitems.need_result_type = True
-
+    _ll_1_odictiter_next = rordereddict._ll_dictnext
     _ll_1_odict_resize = rordereddict.ll_dict_resize
 
     # ---------- strings and unicode ----------

File rpython/rtyper/lltypesystem/rdict.py

                                          ('dict', r_dict.lowleveltype),
                                          ('index', lltype.Signed)))
         self.ll_dictiter = ll_dictiter
-        self.ll_dictnext = ll_dictnext_group[variant]
+        self._ll_dictnext = _ll_dictnext
 
 
 def ll_dictiter(ITERPTR, d):
     iter.index = 0
     return iter
 
-def _make_ll_dictnext(kind):
-    # make three versions of the following function: keys, values, items
-    @jit.look_inside_iff(lambda RETURNTYPE, iter: jit.isvirtual(iter)
-                         and (iter.dict is None or
-                              jit.isvirtual(iter.dict)))
-    @jit.oopspec("dictiter.next%s(iter)" % kind)
-    def ll_dictnext(RETURNTYPE, iter):
-        # note that RETURNTYPE is None for keys and values
-        dict = iter.dict
-        if dict:
-            entries = dict.entries
-            index = iter.index
-            assert index >= 0
-            entries_len = len(entries)
-            while index < entries_len:
-                entry = entries[index]
-                is_valid = entries.valid(index)
-                index = index + 1
-                if is_valid:
-                    iter.index = index
-                    if RETURNTYPE is lltype.Void:
-                        return None
-                    elif kind == 'items':
-                        r = lltype.malloc(RETURNTYPE.TO)
-                        r.item0 = recast(RETURNTYPE.TO.item0, entry.key)
-                        r.item1 = recast(RETURNTYPE.TO.item1, entry.value)
-                        return r
-                    elif kind == 'keys':
-                        return entry.key
-                    elif kind == 'values':
-                        return entry.value
-            # clear the reference to the dict and prevent restarts
-            iter.dict = lltype.nullptr(lltype.typeOf(iter).TO.dict.TO)
-        raise StopIteration
-    return ll_dictnext
-
-ll_dictnext_group = {'keys'  : _make_ll_dictnext('keys'),
-                     'values': _make_ll_dictnext('values'),
-                     'items' : _make_ll_dictnext('items')}
+@jit.look_inside_iff(lambda iter: jit.isvirtual(iter)
+                     and (iter.dict is None or
+                          jit.isvirtual(iter.dict)))
+@jit.oopspec("dictiter.next(iter)")
+def _ll_dictnext(iter):
+    dict = iter.dict
+    if dict:
+        entries = dict.entries
+        index = iter.index
+        assert index >= 0
+        entries_len = len(entries)
+        while index < entries_len:
+            nextindex = index + 1
+            if entries.valid(index):
+                iter.index = nextindex
+                return index
+            index = nextindex
+        # clear the reference to the dict and prevent restarts
+        iter.dict = lltype.nullptr(lltype.typeOf(iter).TO.dict.TO)
+    raise StopIteration
 
 # _____________________________________________________________
 # methods

File rpython/rtyper/lltypesystem/rordereddict.py

         self.variant = variant
         self.lowleveltype = get_ll_dictiter(r_dict.lowleveltype)
         self.ll_dictiter = ll_dictiter
-        self.ll_dictnext = ll_dictnext_group[variant]
+        self._ll_dictnext = _ll_dictnext
 
 
 def ll_dictiter(ITERPTR, d):
     iter.index = 0
     return iter
 
-def _make_ll_dictnext(kind):
-    # make three versions of the following function: keys, values, items
-    @jit.look_inside_iff(lambda RETURNTYPE, iter: jit.isvirtual(iter)
-                         and (iter.dict is None or
-                              jit.isvirtual(iter.dict)))
-    @jit.oopspec("odictiter.next%s(iter)" % kind)
-    def ll_dictnext(RETURNTYPE, iter):
-        # note that RETURNTYPE is None for keys and values
-        dict = iter.dict
-        if not dict:
-            raise StopIteration
-
+@jit.look_inside_iff(lambda iter: jit.isvirtual(iter)
+                     and (iter.dict is None or
+                          jit.isvirtual(iter.dict)))
+@jit.oopspec("odictiter.next(iter)")
+def _ll_dictnext(iter):
+    dict = iter.dict
+    if dict:
         entries = dict.entries
         index = iter.index
         assert index >= 0
         entries_len = dict.num_used_items
         while index < entries_len:
-            entry = entries[index]
-            is_valid = entries.valid(index)
-            index = index + 1
-            if is_valid:
-                iter.index = index
-                if RETURNTYPE is lltype.Void:
-                    return None
-                elif kind == 'items':
-                    r = lltype.malloc(RETURNTYPE.TO)
-                    r.item0 = recast(RETURNTYPE.TO.item0, entry.key)
-                    r.item1 = recast(RETURNTYPE.TO.item1, entry.value)
-                    return r
-                elif kind == 'keys':
-                    return entry.key
-                elif kind == 'values':
-                    return entry.value
-
+            nextindex = index + 1
+            if entries.valid(index):
+                iter.index = nextindex
+                return index
+            index = nextindex
         # clear the reference to the dict and prevent restarts
         iter.dict = lltype.nullptr(lltype.typeOf(iter).TO.dict.TO)
-        raise StopIteration
-
-    return ll_dictnext
-
-ll_dictnext_group = {'keys'  : _make_ll_dictnext('keys'),
-                     'values': _make_ll_dictnext('values'),
-                     'items' : _make_ll_dictnext('items')}
+    raise StopIteration
 
 # _____________________________________________________________
 # methods

File rpython/rtyper/rdict.py

     def rtype_next(self, hop):
         variant = self.variant
         v_iter, = hop.inputargs(self)
-        if variant in ('keys', 'values'):
-            c1 = hop.inputconst(lltype.Void, None)
-        else:
-            c1 = hop.inputconst(lltype.Void, hop.r_result.lowleveltype)
         # record that we know about these two possible exceptions
         hop.has_implicit_exception(StopIteration)
         hop.has_implicit_exception(RuntimeError)
         hop.exception_is_here()
-        v = hop.gendirectcall(self.ll_dictnext, c1, v_iter)
+        v_index = hop.gendirectcall(self._ll_dictnext, v_iter)
+        if variant == 'items' and hop.r_result.lowleveltype != lltype.Void:
+            c1 = hop.inputconst(lltype.Void, hop.r_result.lowleveltype.TO)
+            cflags = hop.inputconst(lltype.Void, {'flavor': 'gc'})
+            v_result = hop.genop('malloc', [c1, cflags],
+                                 resulttype = hop.r_result.lowleveltype)
+        DICT = self.lowleveltype.TO.dict
+        c_dict = hop.inputconst(lltype.Void, 'dict')
+        v_dict = hop.genop('getfield', [v_iter, c_dict], resulttype=DICT)
+        ENTRIES = DICT.TO.entries
+        c_entries = hop.inputconst(lltype.Void, 'entries')
+        v_entries = hop.genop('getfield', [v_dict, c_entries],
+                              resulttype=ENTRIES)
+        if variant != 'values':
+            KEY = ENTRIES.TO.OF.key
+            c_key = hop.inputconst(lltype.Void, 'key')
+            v_key = hop.genop('getinteriorfield', [v_entries, v_index, c_key],
+                              resulttype=KEY)
+        if variant != 'keys':
+            VALUE = ENTRIES.TO.OF.value
+            c_value = hop.inputconst(lltype.Void, 'value')
+            v_value = hop.genop('getinteriorfield', [v_entries,v_index,c_value],
+                                resulttype=VALUE)
         if variant == 'keys':
-            return self.r_dict.recast_key(hop.llops, v)
+            return self.r_dict.recast_key(hop.llops, v_key)
         elif variant == 'values':
-            return self.r_dict.recast_value(hop.llops, v)
+            return self.r_dict.recast_value(hop.llops, v_value)
+        elif hop.r_result.lowleveltype == lltype.Void:
+            return hop.inputconst(lltype.Void, None)
         else:
-            return v
+            assert variant == 'items'
+            ITEM0 = v_result.concretetype.TO.item0
+            ITEM1 = v_result.concretetype.TO.item1
+            if ITEM0 != v_key.concretetype:
+                v_key = hop.genop('cast_pointer', [v_key], resulttype=ITEM0)
+            if ITEM1 != v_value.concretetype:
+                v_value = hop.genop('cast_pointer', [v_value], resulttype=ITEM1)
+            c_item0 = hop.inputconst(lltype.Void, 'item0')
+            c_item1 = hop.inputconst(lltype.Void, 'item1')
+            hop.genop('setfield', [v_result, c_item0, v_key])
+            hop.genop('setfield', [v_result, c_item1, v_value])
+            return v_result

File rpython/rtyper/test/test_rordereddict.py

         rordereddict.ll_dict_setitem(ll_d, llstr("j"), 2)
         ITER = rordereddict.get_ll_dictiter(lltype.Ptr(DICT))
         ll_iter = rordereddict.ll_dictiter(ITER, ll_d)
-        ll_iterkeys = rordereddict.ll_dictnext_group['keys']
-        next = ll_iterkeys(lltype.Signed, ll_iter)
-        assert hlstr(next) == "k"
-        next = ll_iterkeys(lltype.Signed, ll_iter)
-        assert hlstr(next) == "j"
-        py.test.raises(StopIteration, ll_iterkeys, lltype.Signed, ll_iter)
+        ll_dictnext = rordereddict._ll_dictnext
+        num = ll_dictnext(ll_iter)
+        assert hlstr(ll_d.entries[num].key) == "k"
+        num = ll_dictnext(ll_iter)
+        assert hlstr(ll_d.entries[num].key) == "j"
+        py.test.raises(StopIteration, ll_dictnext, ll_iter)
 
     def test_popitem(self):
         DICT = self._get_str_dict()