[mypyc] Faster dict iteration (#8725)

Fixes mypyc/mypyc#167

The implementation is pretty straightforward and follows the idea proposed in the issue. The perf impact is actually pretty small, around 1%

Note: we don't apply the fast path to subclasses here to make dict subclassing safe.

Author: ilevkivskyi

mypyc/ir/rtypes.py

@@ -399,6 +399,16 @@ def deserialize(cls, data: JsonDict, ctx: 'DeserMaps') -> 'RTuple':
 # Exception tuple: (exception class, exception instance, traceback object)
 exc_rtuple = RTuple([object_rprimitive, object_rprimitive, object_rprimitive])
 
+# Dictionary iterator tuple: (should continue, internal offset, key, value)
+# See mypyc.irbuild.for_helpers.ForDictionaryCommon for more details.
+dict_next_rtuple_pair = RTuple(
+    [bool_rprimitive, int_rprimitive, object_rprimitive, object_rprimitive]
+)
+# Same as above but just for key or value.
+dict_next_rtuple_single = RTuple(
+    [bool_rprimitive, int_rprimitive, object_rprimitive]
+)
+
 
 class RInstance(RType):
     """Instance of user-defined class (compiled to C extension class).

mypyc/irbuild/builder.py

@@ -24,6 +24,7 @@
 from mypy.types import (
     Type, Instance, TupleType, UninhabitedType, get_proper_type
 )
+from mypy.maptype import map_instance_to_supertype
 from mypy.visitor import ExpressionVisitor, StatementVisitor
 from mypy.util import split_target
 
@@ -604,6 +605,30 @@ def get_sequence_type(self, expr: Expression) -> RType:
         else:
             return self.type_to_rtype(target_type.args[0])
 
+    def get_dict_base_type(self, expr: Expression) -> Instance:
+        """Find dict type of a dict-like expression.
+
+        This is useful for dict subclasses like SymbolTable.
+        """
+        target_type = get_proper_type(self.types[expr])
+        assert isinstance(target_type, Instance)
+        dict_base = next(base for base in target_type.type.mro
+                         if base.fullname == 'builtins.dict')
+        return map_instance_to_supertype(target_type, dict_base)
+
+    def get_dict_key_type(self, expr: Expression) -> RType:
+        dict_base_type = self.get_dict_base_type(expr)
+        return self.type_to_rtype(dict_base_type.args[0])
+
+    def get_dict_value_type(self, expr: Expression) -> RType:
+        dict_base_type = self.get_dict_base_type(expr)
+        return self.type_to_rtype(dict_base_type.args[1])
+
+    def get_dict_item_type(self, expr: Expression) -> RType:
+        key_type = self.get_dict_key_type(expr)
+        value_type = self.get_dict_value_type(expr)
+        return RTuple([key_type, value_type])
+
     def _analyze_iterable_item_type(self, expr: Expression) -> Type:
         """Return the item type given by 'expr' in an iterable context."""
         # This logic is copied from mypy's TypeChecker.analyze_iterable_item_type.

mypyc/irbuild/for_helpers.py

@@ -6,13 +6,22 @@
 """
 
 from typing import Union, List, Optional, Tuple, Callable
+from typing_extensions import Type, ClassVar
 
-from mypy.nodes import Lvalue, Expression, TupleExpr, CallExpr, RefExpr, GeneratorExpr, ARG_POS
+from mypy.nodes import (
+    Lvalue, Expression, TupleExpr, CallExpr, RefExpr, GeneratorExpr, ARG_POS, MemberExpr
+)
 from mypyc.ir.ops import (
-    Value, BasicBlock, LoadInt, Branch, Register, AssignmentTarget
+    Value, BasicBlock, LoadInt, Branch, Register, AssignmentTarget, TupleGet,
+    AssignmentTargetTuple, TupleSet, OpDescription
 )
 from mypyc.ir.rtypes import (
-    RType, is_short_int_rprimitive, is_list_rprimitive, is_sequence_rprimitive
+    RType, is_short_int_rprimitive, is_list_rprimitive, is_sequence_rprimitive,
+    RTuple, is_dict_rprimitive
+)
+from mypyc.primitives.dict_ops import (
+    dict_next_key_op, dict_next_value_op, dict_next_item_op, dict_check_size_op,
+    dict_key_iter_op, dict_value_iter_op, dict_item_iter_op
 )
 from mypyc.primitives.int_ops import unsafe_short_add
 from mypyc.primitives.list_ops import new_list_op, list_append_op, list_get_item_unsafe_op
@@ -170,6 +179,15 @@ def make_for_loop_generator(builder: IRBuilder,
         for_list.init(expr_reg, target_type, reverse=False)
         return for_list
 
+    if is_dict_rprimitive(rtyp):
+        # Special case "for k in <dict>".
+        expr_reg = builder.accept(expr)
+        target_type = builder.get_dict_key_type(expr)
+
+        for_dict = ForDictionaryKeys(builder, index, body_block, loop_exit, line, nested)
+        for_dict.init(expr_reg, target_type)
+        return for_dict
+
     if (isinstance(expr, CallExpr)
             and isinstance(expr.callee, RefExpr)):
         if (expr.callee.fullname == 'builtins.range'
@@ -233,6 +251,27 @@ def make_for_loop_generator(builder: IRBuilder,
             for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
             for_list.init(expr_reg, target_type, reverse=True)
             return for_list
+    if (isinstance(expr, CallExpr)
+            and isinstance(expr.callee, MemberExpr)
+            and not expr.args):
+        # Special cases for dictionary iterator methods, like dict.items().
+        rtype = builder.node_type(expr.callee.expr)
+        if (is_dict_rprimitive(rtype)
+                and expr.callee.name in ('keys', 'values', 'items')):
+            expr_reg = builder.accept(expr.callee.expr)
+            for_dict_type = None  # type: Optional[Type[ForGenerator]]
+            if expr.callee.name == 'keys':
+                target_type = builder.get_dict_key_type(expr.callee.expr)
+                for_dict_type = ForDictionaryKeys
+            elif expr.callee.name == 'values':
+                target_type = builder.get_dict_value_type(expr.callee.expr)
+                for_dict_type = ForDictionaryValues
+            else:
+                target_type = builder.get_dict_item_type(expr.callee.expr)
+                for_dict_type = ForDictionaryItems
+            for_dict_gen = for_dict_type(builder, index, body_block, loop_exit, line, nested)
+            for_dict_gen.init(expr_reg, target_type)
+            return for_dict_gen
 
     # Default to a generic for loop.
     expr_reg = builder.accept(expr)
@@ -292,6 +331,14 @@ def gen_step(self) -> None:
     def gen_cleanup(self) -> None:
         """Generate post-loop cleanup (if needed)."""
 
+    def load_len(self, expr: Union[Value, AssignmentTarget]) -> Value:
+        """A helper to get collection length, used by several subclasses."""
+        return self.builder.builder.builtin_call(
+            [self.builder.read(expr, self.line)],
+            'builtins.len',
+            self.line,
+        )
+
 
 class ForIterable(ForGenerator):
     """Generate IR for a for loop over an arbitrary iterable (the normal case)."""
@@ -371,17 +418,11 @@ def init(self, expr_reg: Value, target_type: RType, reverse: bool) -> None:
         if not reverse:
             index_reg = builder.add(LoadInt(0))
         else:
-            index_reg = builder.binary_op(self.load_len(), builder.add(LoadInt(1)), '-', self.line)
+            index_reg = builder.binary_op(self.load_len(self.expr_target),
+                                          builder.add(LoadInt(1)), '-', self.line)
         self.index_target = builder.maybe_spill_assignable(index_reg)
         self.target_type = target_type
 
-    def load_len(self) -> Value:
-        return self.builder.builder.builtin_call(
-            [self.builder.read(self.expr_target, self.line)],
-            'builtins.len',
-            self.line,
-        )
-
     def gen_condition(self) -> None:
         builder = self.builder
         line = self.line
@@ -398,7 +439,7 @@ def gen_condition(self) -> None:
             builder.activate_block(second_check)
         # For compatibility with python semantics we recalculate the length
         # at every iteration.
-        len_reg = self.load_len()
+        len_reg = self.load_len(self.expr_target)
         comparison = builder.binary_op(builder.read(self.index_target, line), len_reg, '<', line)
         builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
 
@@ -430,6 +471,136 @@ def gen_step(self) -> None:
              builder.add(LoadInt(step))], line), line)
 
 
+class ForDictionaryCommon(ForGenerator):
+    """Generate optimized IR for a for loop over dictionary keys/values.
+
+    The logic is pretty straightforward, we use PyDict_Next() API wrapped in
+    a tuple, so that we can modify only a single register. The layout of the tuple:
+      * f0: are there more items (bool)
+      * f1: current offset (int)
+      * f2: next key (object)
+      * f3: next value (object)
+    For more info see https://docs.python.org/3/c-api/dict.html#c.PyDict_Next.
+
+    Note that for subclasses we fall back to generic PyObject_GetIter() logic,
+    since they may override some iteration methods in subtly incompatible manner.
+    The fallback logic is implemented in CPy.h via dynamic type check.
+    """
+    dict_next_op = None  # type: ClassVar[OpDescription]
+    dict_iter_op = None  # type: ClassVar[OpDescription]
+
+    def need_cleanup(self) -> bool:
+        # Technically, a dict subclass can raise an unrelated exception
+        # in __next__(), so we need this.
+        return True
+
+    def init(self, expr_reg: Value, target_type: RType) -> None:
+        builder = self.builder
+        self.target_type = target_type
+
+        # We add some variables to environment class, so they can be read across yield.
+        self.expr_target = builder.maybe_spill(expr_reg)
+        offset_reg = builder.add(LoadInt(0))
+        self.offset_target = builder.maybe_spill_assignable(offset_reg)
+        self.size = builder.maybe_spill(self.load_len(self.expr_target))
+
+        # For dict class (not a subclass) this is the dictionary itself.
+        iter_reg = builder.primitive_op(self.dict_iter_op, [expr_reg], self.line)
+        self.iter_target = builder.maybe_spill(iter_reg)
+
+    def gen_condition(self) -> None:
+        """Get next key/value pair, set new offset, and check if we should continue."""
+        builder = self.builder
+        line = self.line
+        self.next_tuple = self.builder.primitive_op(
+            self.dict_next_op, [builder.read(self.iter_target, line),
+                                builder.read(self.offset_target, line)], line)
+
+        # Do this here instead of in gen_step() to minimize variables in environment.
+        new_offset = builder.add(TupleGet(self.next_tuple, 1, line))
+        builder.assign(self.offset_target, new_offset, line)
+
+        should_continue = builder.add(TupleGet(self.next_tuple, 0, line))
+        builder.add(
+            Branch(should_continue, self.body_block, self.loop_exit, Branch.BOOL_EXPR)
+        )
+
+    def gen_step(self) -> None:
+        """Check that dictionary didn't change size during iteration.
+
+        Raise RuntimeError if it is not the case to match CPython behavior.
+        """
+        builder = self.builder
+        line = self.line
+        # Technically, we don't need a new primitive for this, but it is simpler.
+        builder.primitive_op(dict_check_size_op,
+                             [builder.read(self.expr_target, line),
+                              builder.read(self.size, line)], line)
+
+    def gen_cleanup(self) -> None:
+        # Same as for generic ForIterable.
+        self.builder.primitive_op(no_err_occurred_op, [], self.line)
+
+
+class ForDictionaryKeys(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary keys."""
+    dict_next_op = dict_next_key_op
+    dict_iter_op = dict_key_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        # Key is stored at the third place in the tuple.
+        key = builder.add(TupleGet(self.next_tuple, 2, line))
+        builder.assign(builder.get_assignment_target(self.index),
+                       builder.coerce(key, self.target_type, line), line)
+
+
+class ForDictionaryValues(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary values."""
+    dict_next_op = dict_next_value_op
+    dict_iter_op = dict_value_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        # Value is stored at the third place in the tuple.
+        value = builder.add(TupleGet(self.next_tuple, 2, line))
+        builder.assign(builder.get_assignment_target(self.index),
+                       builder.coerce(value, self.target_type, line), line)
+
+
+class ForDictionaryItems(ForDictionaryCommon):
+    """Generate optimized IR for a for loop over dictionary items."""
+    dict_next_op = dict_next_item_op
+    dict_iter_op = dict_item_iter_op
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+
+        key = builder.add(TupleGet(self.next_tuple, 2, line))
+        value = builder.add(TupleGet(self.next_tuple, 3, line))
+
+        # Coerce just in case e.g. key is itself a tuple to be unpacked.
+        assert isinstance(self.target_type, RTuple)
+        key = builder.coerce(key, self.target_type.types[0], line)
+        value = builder.coerce(value, self.target_type.types[1], line)
+
+        target = builder.get_assignment_target(self.index)
+        if isinstance(target, AssignmentTargetTuple):
+            # Simpler code for common case: for k, v in d.items().
+            if len(target.items) != 2:
+                builder.error("Expected a pair for dict item iteration", line)
+            builder.assign(target.items[0], key, line)
+            builder.assign(target.items[1], value, line)
+        else:
+            rvalue = builder.add(TupleSet([key, value], line))
+            builder.assign(target, rvalue, line)
+
+
 class ForRange(ForGenerator):
     """Generate optimized IR for a for loop over an integer range."""