Ejemplo n.º 1
0
    def _compare(self, e1: Exp, e2: Exp, context: Context):
        e1_constant = not free_vars(e1) and not free_funcs(e1)
        e2_constant = not free_vars(e2) and not free_funcs(e2)
        if e1_constant and e2_constant:
            e1v = eval(e1, {})
            e2v = eval(e2, {})
            event("comparison obvious on constants: {} vs {}".format(e1v, e2v))
            return order_objects(e1v, e2v)
        if alpha_equivalent(e1, e2):
            event("shortcutting comparison of identical terms")
            return Order.EQUAL

        path_condition = EAll(context.path_conditions())
        always_le = self.solver.valid(EImplies(path_condition, ELe(e1, e2)))
        always_ge = self.solver.valid(EImplies(path_condition, EGe(e1, e2)))

        if always_le and always_ge:
            return Order.EQUAL
        if always_le:
            return Order.LT
        if always_ge:
            return Order.GT
        return Order.AMBIGUOUS
Ejemplo n.º 2
0
def eval_bulk(
        e : Exp,
        envs : [{str:object}],
        use_default_values_for_undefined_vars : bool = False):
    """Evaluate an expression in many different environments.

    This function accepts the same arguments as `eval`, but takes a list of
    environments instead of just one.

    The call

        eval_bulk(e, envs)

    is equivalent to

        [eval(e, env) for env in envs].

    However, using `eval_bulk` is much faster than repeatedly calling `eval` on
    the same expression.
    """

    if not envs:
        return []

    e = purify(e)
    ops = []
    types = { v.id : v.type for v in free_vars(e) }
    vars = OrderedSet(itertools.chain(types.keys(), free_funcs(e).keys()))
    vmap = { v : i for (i, v) in enumerate(vars) }

    try:
        envs = [ [(env.get(v, mkval(types[v])) if (use_default_values_for_undefined_vars and v in types) else env[v]) for v in vars] for env in envs ]
    except KeyError:
        import sys
        print("OH NO", file=sys.stderr)
        print("e = {}".format(pprint(e)), file=sys.stderr)
        print("eval_bulk({!r}, {!r}, use_default_values_for_undefined_vars={!r})".format(e, envs, use_default_values_for_undefined_vars), file=sys.stderr)
        raise
    _compile(e, vmap, ops)
    return [_eval_compiled(ops, env) for env in envs]
Ejemplo n.º 3
0
def _try_optimize(e : Exp, context : Context, pool : Pool):
    if not accelerate.value:
        return

    if pool != RUNTIME_POOL:
        return

    state_vars = [v for v, p in context.vars() if p == STATE_POOL]
    args = [v for v, p in context.vars() if p == RUNTIME_POOL]

    # ---------------------------------------------------------------------
    # "Rewrite schemes": these trigger on many different AST shapes
    # They are listed first because they are more powerful than the
    # specific rewrite rules below.

    if not free_vars(e) and not free_funcs(e):
        try:
            yield _check(uneval(e.type, eval(e, {})), context, RUNTIME_POOL)
        except NotImplementedError:
            print("Unable to evaluate {!r}".format(e))

    if all(v in state_vars for v in free_vars(e)):
        nsv = strip_EStateVar(e)
        sv = EStateVar(nsv).with_type(e.type)
        yield _check(sv, context, RUNTIME_POOL)

    for ee in fold_into_map(e, context):
        yield _check(ee, context, pool)

    # ---------------------------------------------------------------------
    # "Rewrites": these trigger on specific AST nodes

    if isinstance(e, EBinOp):

        if e.op == "-" and is_collection(e.type):
            ee = optimized_bag_difference(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == "===" and isinstance(e.e1.type, THandle):
            yield _check(EAll([
                optimized_eq(optimized_addr(e.e1), optimized_addr(e.e2)),
                optimized_eq(optimized_val(e.e1),  optimized_val(e.e2)).with_type(BOOL)]), context, RUNTIME_POOL)

        if e.op == BOp.In:
            ee = optimized_in(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, ECond):
        yield _check(optimized_cond(e.cond, e.then_branch, e.else_branch), context, RUNTIME_POOL)

    if isinstance(e, EGetField):
        for ee in optimized_get_field(e.e, e.field_name, args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EListGet) and e.index == ZERO:
        for res in optimized_the(e.e, args):
            yield _check(res, context, RUNTIME_POOL)

    if isinstance(e, EListGet) and isinstance(e.e, ECond):
        yield optimized_cond(e.e.cond,
                             EListGet(e.e.then_branch, e.index).with_type(e.type),
                             EListGet(e.e.else_branch, e.index).with_type(e.type))

    from cozy.structures.treemultiset import ETreeMultisetElems, ETreeMultisetPeek
    if isinstance(e, EListGet) and isinstance(e.e, ETreeMultisetElems):
        yield ETreeMultisetPeek(e.e.e, e.index).with_type(e.type)

    if isinstance(e, EMapGet):
        ee = inline_mapget(e, context)
        yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EUnaryOp):

        if e.op == UOp.Sum:
            for ee in optimized_sum(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Length:
            ee = optimized_len(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Empty:
            ee = optimized_empty(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Exists:
            ee = optimized_exists(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Distinct:
            for ee in optimized_distinct(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.The:
            for ee in optimized_the(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EArgMin) or isinstance(e, EArgMax):
        for ee in optimized_best(e.e, e.key_function, "<" if isinstance(e, EArgMin) else ">", args=args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EFilter):
        for ee in optimized_filter(e.e, e.predicate, args=args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EMap):
        for ee in optimized_map(e.e, e.transform_function, args=args):
            yield _check(ee, context, RUNTIME_POOL)
    from cozy.syntax import ESorted
    from cozy.structures.treemultiset import EMakeMaxTreeMultiset, TMaxTreeMultiset, EMakeMinTreeMultiset, TMinTreeMultiset, ETreeMultisetElems
    target = e
    if isinstance(target, ESorted) and isinstance(target.e, EStateVar):
        e_max = EMakeMaxTreeMultiset(target.e.e).with_type(TMaxTreeMultiset(target.e.e.type.elem_type))
        e_min = EMakeMinTreeMultiset(target.e.e).with_type(TMinTreeMultiset(target.e.e.type.elem_type))
        ee = optimized_cond(target.asc,
                            ETreeMultisetElems(EStateVar(e_min).with_type(e_min.type)).with_type(target.type),
                            ETreeMultisetElems(EStateVar(e_max).with_type(e_max.type)).with_type(target.type))
        yield _check(ee, context, RUNTIME_POOL)
Ejemplo n.º 4
0
def free_vars_and_funcs(e):
    for v in free_vars(e):
        yield v.id
    for f in free_funcs(e):
        yield f
Ejemplo n.º 5
0
def improve(
        target : Exp,
        assumptions : Exp,
        binders : [EVar],
        state_vars : [EVar],
        args : [EVar],
        cost_model : CostModel,
        builder : ExpBuilder,
        stop_callback = never_stop,
        hints : [Exp] = None,
        examples = None):
    """
    Improve the target expression using enumerative synthesis.
    This function is a generator that yields increasingly better and better
    versions of the input expression `target`.

    Notes on internals of this algorithm follow.

    Key differences from "regular" enumerative synthesis:
        - Expressions may be built using a set of "binders"---extra free
          variables thrown into the mix at the beginning.
        - Expressions are either "state" expressions or "runtime" expressions,
          allowing this algorithm to choose what things to store on the data
          structure and what things to compute at query execution time. (The
          cost model is ultimately responsible for this choice.)

    Other features of this algorithm:
        - If a better version of *any subexpression* for the target is found,
          it is immediately substituted in and the overall expression is
          returned. This "smooths out" the search space a little, and lets us
          find kinda-good solutions very quickly, even if the best possible
          solution is out of reach.
    """

    print("call to improve:")
    print("""improve(
        target={target!r},
        assumptions={assumptions!r},
        binders={binders!r},
        state_vars={state_vars!r},
        args={args!r},
        cost_model={cost_model!r},
        builder={builder!r},
        stop_callback={stop_callback!r},
        hints={hints!r},
        examples={examples!r})""".format(
            target=target,
            assumptions=assumptions,
            binders=binders,
            state_vars=state_vars,
            args=args,
            cost_model=cost_model,
            builder=builder,
            stop_callback=stop_callback,
            hints=hints,
            examples=examples))

    print()
    print("improving: {}".format(pprint(target)))
    print("subject to: {}".format(pprint(assumptions)))
    print()

    assert exp_wf(
        target,
        state_vars=set(state_vars),
        args=set(args),
        assumptions=assumptions)

    binders = list(binders)
    target = fixup_binders(target, binders, allow_add=False)
    hints = [fixup_binders(h, binders, allow_add=False) for h in (hints or ())]
    assumptions = fixup_binders(assumptions, binders, allow_add=False)
    builder = FixedBuilder(builder, state_vars, args, binders, assumptions)
    target_cost = cost_model.cost(target, RUNTIME_POOL)

    if eliminate_vars.value and can_elim_vars(target, assumptions, state_vars):
        print("This job does not depend on state_vars.")
        builder = StateElimBuilder(builder)

    vars = list(free_vars(target) | free_vars(assumptions))
    funcs = free_funcs(EAll([target, assumptions]))

    solver = None
    if incremental.value:
        solver = IncrementalSolver(vars=vars, funcs=funcs, collection_depth=check_depth.value)
        solver.add_assumption(assumptions)
        _sat = solver.satisfy
    else:
        _sat = lambda e: satisfy(e, vars=vars, funcs=funcs, collection_depth=check_depth.value)

    if _sat(T) is None:
        print("assumptions are unsat; this query will never be called")
        yield construct_value(target.type)
        return

    if examples is None:
        examples = []
    learner = Learner(target, assumptions, binders, state_vars, args, vars + binders, examples, cost_model, builder, stop_callback, hints, solver=solver)
    try:
        while True:
            # 1. find any potential improvement to any sub-exp of target
            try:
                old_e, new_e, local_assumptions, repl = learner.next()
            except NoMoreImprovements:
                break

            # 2. substitute-in the improvement
            print("Found candidate replacement [{}] for [{}] in".format(pprint(new_e), pprint(old_e)))
            print(pprint(repl(EVar("@___"))))
            new_target = repl(new_e)

            # 3. check
            if incremental.value:
                solver.push()
                solver.add_assumption(ENot(EBinOp(target, "==", new_target).with_type(BOOL)))
                counterexample = _sat(T)
            else:
                formula = EAll([assumptions, ENot(EBinOp(target, "==", new_target).with_type(BOOL))])
                counterexample = _sat(formula)
            if counterexample is not None:

                # Ok they aren't equal.  Now we need an example that
                # differentiates BOTH target/new_target AND old_e/new_e.
                if incremental.value:
                    counterexample = _sat(EAll([
                            EAll(local_assumptions),
                            ENot(EBinOp(old_e,  "===", new_e).with_type(BOOL))]))
                else:
                    counterexample = _sat(EAll([
                            assumptions,
                            EAll(local_assumptions),
                            ENot(EBinOp(target, "==", new_target).with_type(BOOL)),
                            ENot(EBinOp(old_e,  "===", new_e).with_type(BOOL))]))
                if counterexample is None:
                    print("!!! unable to satisfy top- and sub-expressions")
                    print("assumptions = {!r}".format(assumptions))
                    print("local_assumptions = {!r}".format(EAll(local_assumptions)))
                    print("old_e = {!r}".format(old_e))
                    print("target = {!r}".format(target))
                    print("new_e = {!r}".format(new_e))
                    print("new_target = {!r}".format(new_target))
                    raise Exception("unable to find an example that differentiates both the toplevel- and sub-expressions")

                if counterexample in examples:
                    print("assumptions = {!r}".format(assumptions))
                    print("duplicate example: {!r}".format(counterexample))
                    print("old target = {!r}".format(target))
                    print("new target = {!r}".format(new_target))
                    print("old fp = {}".format(learner._fingerprint(old_e)))
                    print("new fp = {}".format(learner._fingerprint(new_e)))
                    print("old target fp = {}".format(learner._fingerprint(target)))
                    print("new target fp = {}".format(learner._fingerprint(new_target)))
                    raise Exception("got a duplicate example")
                # a. if incorrect: add example, reset the learner
                examples.append(counterexample)
                print("new example: {}".format(truncate(repr(counterexample))))
                print("restarting with {} examples".format(len(examples)))
                learner.reset(examples)
            else:
                # b. if correct: yield it, watch the new target, goto 1

                if CHECK_FINAL_COST:
                    new_cost = cost_model.cost(new_target, RUNTIME_POOL)
                    print("cost: {} -----> {}".format(target_cost, new_cost))
                    if incremental.value:
                        ordering = new_cost.compare_to(target_cost, solver=solver)
                    else:
                        ordering = new_cost.compare_to(target_cost, assumptions=assumptions)
                    if ordering == Cost.WORSE:
                        if CHECK_SUBST_COST:
                            print("WHOOPS! COST GOT WORSE!")
                            if save_testcases.value:
                                with open(save_testcases.value, "a") as f:
                                    f.write("def testcase():\n")
                                    f.write("    costmodel = {}\n".format(repr(cost_model)))
                                    f.write("    old_e = {}\n".format(repr(old_e)))
                                    f.write("    new_e = {}\n".format(repr(new_e)))
                                    f.write("    target = {}\n".format(repr(target)))
                                    f.write("    new_target = {}\n".format(repr(new_target)))
                                    f.write("    if costmodel.cost(new_e, RUNTIME_POOL) <= costmodel.cost(old_e, RUNTIME_POOL) and costmodel.cost(new_target, RUNTIME_POOL) > costmodel.cost(target, RUNTIME_POOL):\n")
                                    f.write('        for name, x in zip(["old_e", "new_e", "target", "new_target"], [old_e, new_e, target, new_target]):\n')
                                    f.write('            print("{}: {}".format(name, pprint(x)))\n')
                                    f.write('            print("    cost = {}".format(costmodel.cost(x, RUNTIME_POOL)))\n')
                                    f.write("        assert False\n")
                            # raise Exception("detected nonmonotonicity")
                        else:
                            print("*** cost is worse")
                            # print(repr(target))
                            # print(repr(new_target))
                        continue
                    elif ordering == Cost.UNORDERED:
                        print("*** cost is unchanged")
                        # print(repr(target))
                        # print(repr(new_target))
                        continue
                    target_cost = new_cost
                print("found improvement: {} -----> {}".format(pprint(old_e), pprint(new_e)))
                # print(repr(target))
                # print(repr(new_target))

                # binders are not allowed to "leak" out
                to_yield = new_target
                if any(v in binders for v in free_vars(new_target)):
                    print("WARNING: stripping binders in {}".format(pprint(new_target)), file=sys.stderr)
                    to_yield = subst(new_target, { b.id : construct_value(b.type) for b in binders })
                yield to_yield

                if reset_on_success.value and (not CHECK_FINAL_COST or ordering != Cost.UNORDERED):
                    learner.reset(examples)
                learner.watch(new_target)
                target = new_target

                if heuristic_done(new_target, args):
                    print("target now matches doneness heuristic")
                    break
            if incremental.value:
                solver.pop()

    except KeyboardInterrupt:
        for e in learner.cache.random_sample(50):
            print(pprint(e))
        raise
Ejemplo n.º 6
0
def _try_optimize(e: Exp, context: Context, pool: Pool):
    if not accelerate.value:
        return

    if pool != RUNTIME_POOL:
        return

    state_vars = [v for v, p in context.vars() if p == STATE_POOL]
    args = [v for v, p in context.vars() if p == RUNTIME_POOL]

    # ---------------------------------------------------------------------
    # "Rewrite schemes": these trigger on many different AST shapes
    # They are listed first because they are more powerful than the
    # specific rewrite rules below.

    if not free_vars(e) and not free_funcs(e):
        try:
            yield _check(uneval(e.type, eval(e, {})), context, RUNTIME_POOL)
        except NotImplementedError:
            print("Unable to evaluate {!r}".format(e))

    if all(v in state_vars for v in free_vars(e)):
        nsv = strip_EStateVar(e)
        sv = EStateVar(nsv).with_type(e.type)
        yield _check(sv, context, RUNTIME_POOL)

    for ee in fold_into_map(e, context):
        yield _check(ee, context, pool)

    # ---------------------------------------------------------------------
    # "Rewrites": these trigger on specific AST nodes

    if isinstance(e, EBinOp):

        if e.op == "-" and is_collection(e.type):
            ee = optimized_bag_difference(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == "===" and isinstance(e.e1.type, THandle):
            yield _check(
                EAll([
                    optimized_eq(optimized_addr(e.e1), optimized_addr(e.e2)),
                    optimized_eq(optimized_val(e.e1),
                                 optimized_val(e.e2)).with_type(BOOL)
                ]), context, RUNTIME_POOL)

        if e.op == BOp.In:
            ee = optimized_in(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, ECond):
        yield _check(optimized_cond(e.cond, e.then_branch, e.else_branch),
                     context, RUNTIME_POOL)

    if isinstance(e, EGetField):
        for ee in optimized_get_field(e.e, e.field_name, args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EListGet) and e.index == ZERO:
        for res in optimized_the(e.e, args):
            yield _check(res, context, RUNTIME_POOL)

    if isinstance(e, EListGet) and isinstance(e.e, ECond):
        yield optimized_cond(
            e.e.cond,
            EListGet(e.e.then_branch, e.index).with_type(e.type),
            EListGet(e.e.else_branch, e.index).with_type(e.type))

    from cozy.structures.treemultiset import ETreeMultisetElems, ETreeMultisetPeek
    if isinstance(e, EListGet) and isinstance(e.e, ETreeMultisetElems):
        yield ETreeMultisetPeek(e.e.e, e.index).with_type(e.type)

    if isinstance(e, EMapGet):
        ee = inline_mapget(e, context)
        yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EUnaryOp):

        if e.op == UOp.Sum:
            for ee in optimized_sum(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Length:
            ee = optimized_len(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Empty:
            ee = optimized_empty(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Exists:
            ee = optimized_exists(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.Distinct:
            for ee in optimized_distinct(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if e.op == UOp.The:
            for ee in optimized_the(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EArgMin) or isinstance(e, EArgMax):
        for ee in optimized_best(e.e,
                                 e.key_function,
                                 "<" if isinstance(e, EArgMin) else ">",
                                 args=args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EFilter):
        for ee in optimized_filter(e.e, e.predicate, args=args):
            yield _check(ee, context, RUNTIME_POOL)

    if isinstance(e, EMap):
        for ee in optimized_map(e.e, e.transform_function, args=args):
            yield _check(ee, context, RUNTIME_POOL)
    from cozy.syntax import ESorted
    from cozy.structures.treemultiset import EMakeMaxTreeMultiset, TMaxTreeMultiset, EMakeMinTreeMultiset, TMinTreeMultiset, ETreeMultisetElems
    target = e
    if isinstance(target, ESorted) and isinstance(target.e, EStateVar):
        e_max = EMakeMaxTreeMultiset(target.e.e).with_type(
            TMaxTreeMultiset(target.e.e.type.elem_type))
        e_min = EMakeMinTreeMultiset(target.e.e).with_type(
            TMinTreeMultiset(target.e.e.type.elem_type))
        ee = optimized_cond(
            target.asc,
            ETreeMultisetElems(EStateVar(e_min).with_type(
                e_min.type)).with_type(target.type),
            ETreeMultisetElems(EStateVar(e_max).with_type(
                e_max.type)).with_type(target.type))
        yield _check(ee, context, RUNTIME_POOL)
Ejemplo n.º 7
0
def _try_optimize(e, context, pool):
    if not accelerate.value:
        return

    state_vars = [v for v, p in context.vars() if p == STATE_POOL]
    args = [v for v, p in context.vars() if p == RUNTIME_POOL]

    if pool == RUNTIME_POOL:

        if not free_vars(e) and not free_funcs(e):
            try:
                yield _check(uneval(e.type, eval(e, {})), context,
                             RUNTIME_POOL)
            except NotImplementedError:
                print("Unable to evaluate {!r}".format(e))

        if all(v in state_vars for v in free_vars(e)):
            nsv = strip_EStateVar(e)
            sv = EStateVar(nsv).with_type(e.type)
            yield _check(sv, context, RUNTIME_POOL)

        for ee in fold_into_map(e, context):
            yield _check(ee, context, pool)

        if isinstance(e, EListGet) and e.index == ZERO:
            for res in optimize_the(e.e, args):
                yield _check(res, context, RUNTIME_POOL)

        if isinstance(e, EArgMin) or isinstance(e, EArgMax):
            for ee in optimized_best(e.e,
                                     e.f,
                                     "<" if isinstance(e, EArgMin) else ">",
                                     args=args):
                yield _check(ee, context, RUNTIME_POOL)

        if is_collection(e.type) and isinstance(e, EBinOp) and e.op == "-":
            ee = optimized_bag_difference(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EBinOp) and e.op == "===" and isinstance(
                e.e1.type, THandle):
            yield _check(
                EAll([
                    optimized_eq(optimized_addr(e.e1), optimized_addr(e.e2)),
                    optimized_eq(optimized_val(e.e1),
                                 optimized_val(e.e2)).with_type(BOOL)
                ]), context, RUNTIME_POOL)

        if isinstance(e, ECond):
            yield _check(optimized_cond(e.cond, e.then_branch, e.else_branch),
                         context, RUNTIME_POOL)

        if isinstance(e, EGetField):
            for ee in optimized_get_field(e.e, e.f, args):
                yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EBinOp) and e.op == BOp.In:
            ee = optimized_in(e.e1, e.e2)
            yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EUnaryOp) and e.op == UOp.Sum:
            for ee in optimized_sum(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EUnaryOp) and e.op == UOp.Empty:
            ee = optimized_empty(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EUnaryOp) and e.op == UOp.Exists:
            ee = optimized_exists(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EUnaryOp) and e.op == UOp.Length:
            ee = optimized_len(e.e)
            yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EUnaryOp) and e.op == UOp.The:
            for ee in optimize_the(e.e, args):
                yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EFilter):
            ee = optimize_filter_as_if_distinct(e.e, e.p, args=args)
            yield _check(ee, context, RUNTIME_POOL)
            if isinstance(e.e, EFilter):
                # try swizzle
                ee = EFilter(_simple_filter(e.e.e, e.p, args=args),
                             e.e.p).with_type(e.type)
                yield _check(ee, context, RUNTIME_POOL)

        if isinstance(e, EMap):
            for ee in optimize_map(e.e, e.f, args=args):
                yield _check(ee, context, RUNTIME_POOL)