def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expression, body_insts: Callable[[Value], None], line: int) -> None: """Generate IR for a sequence iteration. This function only works for sequence type. Compared to for_loop_helper, it would feed iteration index to body_insts. Args: index: the loop index Lvalue expr: the expression to iterate over body_insts: a function that generates the body of the loop. It needs a index as parameter. """ expr_reg = builder.accept(expr) assert is_sequence_rprimitive(expr_reg.type) target_type = builder.get_sequence_type(expr) body_block = BasicBlock() step_block = BasicBlock() exit_block = BasicBlock() condition_block = BasicBlock() for_gen = ForSequence(builder, index, body_block, exit_block, line, False) for_gen.init(expr_reg, target_type, reverse=False) builder.push_loop_stack(step_block, exit_block) builder.goto_and_activate(condition_block) for_gen.gen_condition() builder.activate_block(body_block) for_gen.begin_body() body_insts(builder.read(for_gen.index_target)) builder.goto_and_activate(step_block) for_gen.gen_step() builder.goto(condition_block) for_gen.add_cleanup(exit_block) builder.pop_loop_stack() builder.activate_block(exit_block)
def make_for_loop_generator(builder: IRBuilder, index: Lvalue, expr: Expression, body_block: BasicBlock, loop_exit: BasicBlock, line: int, nested: bool = False) -> 'ForGenerator': """Return helper object for generating a for loop over an iterable. If "nested" is True, this is a nested iterator such as "e" in "enumerate(e)". """ rtyp = builder.node_type(expr) if is_sequence_rprimitive(rtyp): # Special case "for x in <list>". expr_reg = builder.accept(expr) target_type = builder.get_sequence_type(expr) for_list = ForSequence(builder, index, body_block, loop_exit, line, nested) 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' and (len(expr.args) <= 2 or (len(expr.args) == 3 and builder.extract_int(expr.args[2]) is not None)) and set(expr.arg_kinds) == {ARG_POS}): # Special case "for x in range(...)". # We support the 3 arg form but only for int literals, since it doesn't # seem worth the hassle of supporting dynamically determining which # direction of comparison to do. if len(expr.args) == 1: start_reg = builder.add(LoadInt(0)) end_reg = builder.accept(expr.args[0]) else: start_reg = builder.accept(expr.args[0]) end_reg = builder.accept(expr.args[1]) if len(expr.args) == 3: step = builder.extract_int(expr.args[2]) assert step is not None if step == 0: builder.error("range() step can't be zero", expr.args[2].line) else: step = 1 for_range = ForRange(builder, index, body_block, loop_exit, line, nested) for_range.init(start_reg, end_reg, step) return for_range elif (expr.callee.fullname == 'builtins.enumerate' and len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(index, TupleExpr) and len(index.items) == 2): # Special case "for i, x in enumerate(y)". lvalue1 = index.items[0] lvalue2 = index.items[1] for_enumerate = ForEnumerate(builder, index, body_block, loop_exit, line, nested) for_enumerate.init(lvalue1, lvalue2, expr.args[0]) return for_enumerate elif (expr.callee.fullname == 'builtins.zip' and len(expr.args) >= 2 and set(expr.arg_kinds) == {ARG_POS} and isinstance(index, TupleExpr) and len(index.items) == len(expr.args)): # Special case "for x, y in zip(a, b)". for_zip = ForZip(builder, index, body_block, loop_exit, line, nested) for_zip.init(index.items, expr.args) return for_zip if (expr.callee.fullname == 'builtins.reversed' and len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and is_sequence_rprimitive(rtyp)): # Special case "for x in reversed(<list>)". expr_reg = builder.accept(expr.args[0]) target_type = builder.get_sequence_type(expr) 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) for_obj = ForIterable(builder, index, body_block, loop_exit, line, nested) item_type = builder._analyze_iterable_item_type(expr) item_rtype = builder.type_to_rtype(item_type) for_obj.init(expr_reg, item_rtype) return for_obj