def enumerate(self, context, size, pool, enumerate_subexps, enumerate_lambdas):
from cozy.typecheck import is_collection
if pool == STATE_POOL:
for (sz1, sz2) in pick_to_sum(2, size-1):
for e in enumerate_subexps(context, sz1, pool):
if is_collection(e.type):
elem_type = e.type.elem_type
yield EMakeMaxTreeMultiset(e).with_type(TMaxTreeMultiset(elem_type))
yield EMakeMinTreeMultiset(e).with_type(TMinTreeMultiset(elem_type))
def enumerate(self, context, size, pool, enumerate_subexps,
enumerate_lambdas):
from cozy.typecheck import is_collection
if pool == STATE_POOL:
for (sz1, sz2) in pick_to_sum(2, size - 1):
for e in enumerate_subexps(context, sz1, pool):
if is_collection(e.type):
elem_type = e.type.elem_type
yield EMakeMaxTreeMultiset(e).with_type(
TMaxTreeMultiset(elem_type))
yield EMakeMinTreeMultiset(e).with_type(
TMinTreeMultiset(elem_type))
def expressions_may_exist_above_size(self, context, pool, size):
"""Returns true if expressions larger than `size` might exist.
If this method returns false, then all calls to `enumerate` or
`enumerate_with_info` with sizes strictly larger than the given size
will yield no results.
"""
if size <= 0:
return True
maximum_arity = 3 # TODO: adjust this later?
for arity in range(1, maximum_arity + 1):
for split in pick_to_sum(arity, size):
for exprs in itertools.product(
*(self.enumerate(context, sz, pool) for sz in split)):
return True
if pool == RUNTIME_POOL:
return self.expressions_may_exist_above_size(
context.root(), STATE_POOL, size - 1)
return False
def enumerate(self, context, size, pool, enumerate_subexps, enumerate_lambdas):
from cozy.typecheck import is_ordered, is_collection
if pool == STATE_POOL:
for (sz1, sz2) in pick_to_sum(2, size-1):
for e in enumerate_subexps(context, sz1, pool):
if is_collection(e.type):
elem_type = e.type.elem_type
for keyfunc in enumerate_lambdas(e, pool, sz2):
key_type = keyfunc.body.type
if is_ordered(key_type):
yield EMakeMinHeap(e, keyfunc).with_type(TMinHeap(elem_type, key_type))
yield EMakeMaxHeap(e, keyfunc).with_type(TMaxHeap(elem_type, key_type))
elif pool == RUNTIME_POOL:
for e in enumerate_subexps(context.root(), size-1, STATE_POOL):
t = e.type
if isinstance(t, TMinHeap) or isinstance(t, TMaxHeap):
elem_type = t.elem_type
# yielding EHeapElems would be redundant
yield EHeapPeek (EStateVar(e).with_type(e.type)).with_type(elem_type)
yield EHeapPeek2(EStateVar(e).with_type(e.type)).with_type(elem_type)
def build(self, cache, size):
# print("Cache:")
# for (e, sz, pool) in cache:
# from cozy.syntax_tools import pprint
# print(" @size={}, pool={}\t:\t{}".format(sz, pool, pprint(e)))
binders_by_type = group_by(self.binders, lambda b: b.type)
for pool in ALL_POOLS:
if size == 1:
yield self.check(T, pool)
yield self.check(F, pool)
yield self.check(ZERO, pool)
yield self.check(ONE, pool)
for b in self.binders:
yield self.check(b, pool)
if pool == STATE_POOL:
for v in self.state_vars:
yield self.check(v, pool)
elif pool == RUNTIME_POOL:
for v in self.args:
yield self.check(v, pool)
if not build_exprs.value:
return
for e in cache.find(pool=STATE_POOL, size=size - 1):
if all(v in self.state_vars for v in free_vars(e)):
yield self.check(
EStateVar(e).with_type(e.type), RUNTIME_POOL)
for e in cache.find(pool=pool, size=size - 1):
t = TBag(e.type)
yield self.check(EEmptyList().with_type(t), pool)
yield self.check(ESingleton(e).with_type(t), pool)
for e in cache.find(pool=pool, type=TRecord, size=size - 1):
for (f, t) in e.type.fields:
yield self.check(EGetField(e, f).with_type(t), pool)
for e in cache.find_collections(pool=pool, size=size - 1):
if is_numeric(e.type.t):
yield self.check(
EUnaryOp(UOp.Sum, e).with_type(e.type.t), pool)
for e in cache.find(pool=pool, type=THandle, size=size - 1):
yield self.check(
EGetField(e, "val").with_type(e.type.value_type), pool)
for e in cache.find(pool=pool, type=TTuple, size=size - 1):
for n in range(len(e.type.ts)):
yield self.check(
ETupleGet(e, n).with_type(e.type.ts[n]), pool)
for e in cache.find(pool=pool, type=BOOL, size=size - 1):
yield self.check(EUnaryOp(UOp.Not, e).with_type(BOOL), pool)
for e in cache.find(pool=pool, type=INT, size=size - 1):
yield self.check(EUnaryOp("-", e).with_type(INT), pool)
for m in cache.find(pool=pool, type=TMap, size=size - 1):
yield self.check(EMapKeys(m).with_type(TBag(m.type.k)), pool)
for (sz1, sz2) in pick_to_sum(2, size - 1):
for a1 in cache.find(pool=pool, size=sz1):
if not is_numeric(a1.type):
continue
for a2 in cache.find(pool=pool, type=a1.type, size=sz2):
yield self.check(
EBinOp(a1, "+", a2).with_type(INT), pool)
yield self.check(
EBinOp(a1, "-", a2).with_type(INT), pool)
yield self.check(
EBinOp(a1, ">", a2).with_type(BOOL), pool)
yield self.check(
EBinOp(a1, "<", a2).with_type(BOOL), pool)
yield self.check(
EBinOp(a1, ">=", a2).with_type(BOOL), pool)
yield self.check(
EBinOp(a1, "<=", a2).with_type(BOOL), pool)
for a1 in cache.find_collections(pool=pool, size=sz1):
for a2 in cache.find(pool=pool, type=a1.type, size=sz2):
yield self.check(
EBinOp(a1, "+", a2).with_type(a1.type), pool)
yield self.check(
EBinOp(a1, "-", a2).with_type(a1.type), pool)
for a2 in cache.find(pool=pool, type=a1.type.t, size=sz2):
yield self.check(
EBinOp(a2, BOp.In, a1).with_type(BOOL), pool)
for a1 in cache.find(pool=pool, type=BOOL, size=sz1):
for a2 in cache.find(pool=pool, type=BOOL, size=sz2):
yield self.check(
EBinOp(a1, BOp.And, a2).with_type(BOOL), pool)
yield self.check(
EBinOp(a1, BOp.Or, a2).with_type(BOOL), pool)
for a1 in cache.find(pool=pool, size=sz1):
if not isinstance(a1.type, TMap):
for a2 in cache.find(pool=pool, type=a1.type,
size=sz2):
yield self.check(EEq(a1, a2), pool)
yield self.check(
EBinOp(a1, "!=", a2).with_type(BOOL), pool)
for m in cache.find(pool=pool, type=TMap, size=sz1):
for k in cache.find(pool=pool, type=m.type.k, size=sz2):
yield self.check(
EMapGet(m, k).with_type(m.type.v), pool)
yield self.check(EHasKey(m, k).with_type(BOOL), pool)
for (sz1, sz2, sz3) in pick_to_sum(3, size - 1):
for cond in cache.find(pool=pool, type=BOOL, size=sz1):
for then_branch in cache.find(pool=pool, size=sz2):
for else_branch in cache.find(pool=pool,
size=sz3,
type=then_branch.type):
yield self.check(
ECond(cond, then_branch,
else_branch).with_type(then_branch.type),
pool)
for bag in cache.find_collections(pool=pool, size=size - 1):
# len of bag
count = EUnaryOp(UOp.Length, bag).with_type(INT)
yield self.check(count, pool)
# empty?
yield self.check(
EUnaryOp(UOp.Empty, bag).with_type(BOOL), pool)
# exists?
yield self.check(
EUnaryOp(UOp.Exists, bag).with_type(BOOL), pool)
# singleton?
yield self.check(EEq(count, ONE), pool)
yield self.check(
EUnaryOp(UOp.The, bag).with_type(bag.type.t), pool)
yield self.check(
EUnaryOp(UOp.Distinct, bag).with_type(bag.type), pool)
yield self.check(
EUnaryOp(UOp.AreUnique, bag).with_type(BOOL), pool)
if bag.type.t == BOOL:
yield self.check(
EUnaryOp(UOp.Any, bag).with_type(BOOL), pool)
yield self.check(
EUnaryOp(UOp.All, bag).with_type(BOOL), pool)
for (sz1, sz2) in pick_to_sum(2, size - 1):
for bag in cache.find_collections(pool=pool, size=sz1):
for binder in binders_by_type[bag.type.t]:
for body in itertools.chain(
cache.find(pool=pool, size=sz2), (binder, )):
yield self.check(
EMap(bag,
ELambda(binder,
body)).with_type(TBag(body.type)),
pool)
if body.type == BOOL:
yield self.check(
EFilter(bag,
ELambda(binder,
body)).with_type(bag.type),
pool)
if body.type == INT:
yield self.check(
EArgMin(bag, ELambda(
binder, body)).with_type(bag.type.t),
pool)
yield self.check(
EArgMax(bag, ELambda(
binder, body)).with_type(bag.type.t),
pool)
if pool == RUNTIME_POOL and isinstance(
body.type, TBag):
yield self.check(
EFlatMap(bag,
ELambda(binder, body)).with_type(
TBag(body.type.t)), pool)
for (sz1, sz2) in pick_to_sum(2, size - 1):
for bag in cache.find_collections(pool=STATE_POOL, size=sz1):
if not is_scalar(bag.type.t):
continue
for b in binders_by_type[bag.type.t]:
for val in cache.find(pool=STATE_POOL, size=sz2):
t = TMap(bag.type.t, val.type)
m = EMakeMap2(bag, ELambda(b, val)).with_type(t)
yield self.check(m, STATE_POOL)
def build(self, cache, size):
for e in cache.find(pool=RUNTIME_POOL, size=size - 1, type=INT):
if not is_root(e):
continue
e2 = simplify_sum(e)
if e != e2:
yield self.check(e2, RUNTIME_POOL)
# for e in cache.find(pool=RUNTIME_POOL, size=size-1):
# if isinstance(e, EMapGet) and isinstance(e.map, EMakeMap2):
# x = e.map.value.apply_to(e.key)
# x._tag = True
# yield self.check(x, RUNTIME_POOL)
# [x] - ys
for e in cache.find_collections(pool=RUNTIME_POOL, size=size - 1):
if not is_root(e):
continue
if isinstance(e, EBinOp) and e.op == "-" and isinstance(
e.e1, ESingleton):
x = e.e1.e
y = e.e2
x = ECond(
EBinOp(x, BOp.In, y).with_type(BOOL),
EEmptyList().with_type(e.type), e.e1).with_type(e.type)
yield self.check(x, RUNTIME_POOL)
elif isinstance(e, EUnaryOp) and e.op == UOp.Distinct:
e = strip_EStateVar(e)
m = EMakeMap2(e.e, mk_lambda(e.type.t, lambda x: T)).with_type(
TMap(e.type.t, BOOL))
yield self.check(m, STATE_POOL)
m = EStateVar(m).with_type(m.type)
yield self.check(m, RUNTIME_POOL)
x = EMapKeys(m).with_type(e.type)
# x._tag = True
yield self.check(x, RUNTIME_POOL)
# # x in ys ----> (count x in ys) > 0
# for e in cache.find(pool=RUNTIME_POOL, type=BOOL, size=size-1):
# if isinstance(e, EBinOp) and e.op == BOp.In:
# for b in self.binders:
# if b.type != e.e1.type:
# continue
# x = EGt(
# EUnaryOp(UOp.Length, EFilter(e.e2, ELambda(b, EEq(e.e1, b))).with_type(e.e2.type)).with_type(INT),
# ZERO)
# x._tag = True
# yield self.check(x, RUNTIME_POOL)
for e in cache.find(pool=RUNTIME_POOL, size=size - 1):
if not is_root(e):
continue
if (isinstance(e, EArgMin) or isinstance(
e, EArgMax)) and isinstance(e.e, EBinOp) and e.e.op == "+":
l = e.e.e1
r = e.e.e2
op = e.e.op
f = lambda x: type(e)(x, e.f).with_type(e.type)
ll = EStateVar(f(l.e)).with_type(e.type) if isinstance(
l, EStateVar) else f(l)
rr = EStateVar(f(r.e)).with_type(e.type) if isinstance(
r, EStateVar) else f(r)
x = ECond(
EUnaryOp(UOp.Exists, l).with_type(BOOL),
ECond(
EUnaryOp(UOp.Exists, r).with_type(BOOL),
f(
EBinOp(
ESingleton(ll).with_type(e.e.type), op,
ESingleton(rr).with_type(e.e.type)).with_type(
e.e.type)), ll).with_type(e.type),
rr).with_type(e.type)
# from cozy.solver import valid
# assert valid(EEq(e, x), model_callback=print)
x._tag = True
yield self.check(x, RUNTIME_POOL)
# is-last(x, l)
for (sz1, sz2) in pick_to_sum(2, size - 1):
for e1 in cache.find(pool=RUNTIME_POOL, size=sz1):
if not is_root(e1):
continue
for e2 in cache.find_collections(pool=STATE_POOL,
size=sz2,
of=e1.type):
if not is_root(e2):
continue
for b in self.binders:
if b.type != e1.type:
continue
m = EMakeMap2(
e2,
mk_lambda(
e2.type.t, lambda x: EUnaryOp(
UOp.Length,
EFilter(
e2,
mk_lambda(e2.type.t, lambda y: EEq(
x, y))).with_type(e2.type)).
with_type(INT))).with_type(TMap(
e2.type.t, INT))
# filt = EFilter(e2, ELambda(b, EEq(e1, b))).with_type(e2.type)
# x = EEq(
# EUnaryOp(UOp.Length, filt).with_type(INT),
# ONE)
x = EGt(
EMapGet(EStateVar(m).with_type(m.type),
e1).with_type(INT), ONE)
# x._tag = True
yield self.check(x, RUNTIME_POOL)
# histogram
# for e in cache.find_collections(pool=STATE_POOL, size=size-1):
# m = EMakeMap2(e,
# mk_lambda(e.type.t, lambda x:
# EUnaryOp(UOp.Length, EFilter(e,
# mk_lambda(e.type.t, lambda y: EEq(x, y))).with_type(e.type)).with_type(INT))).with_type(TMap(e.type.t, INT))
# m._tag = True
# yield self.check(m, STATE_POOL)
# Fixup EFilter(\x -> ECond...)
for e in cache.find_collections(pool=RUNTIME_POOL, size=size - 1):
if not is_root(e):
continue
if isinstance(e, EFilter):
for (_, x, r, _) in enumerate_fragments(e.p.body):
if isinstance(x, ECond):
lhs = EFilter(
e.e,
ELambda(e.p.arg, EAll([x.cond,
r(x.then_branch)
]))).with_type(e.type)
rhs = EFilter(
e.e,
ELambda(e.p.arg,
EAll([ENot(x.cond),
r(x.else_branch)
]))).with_type(e.type)
union = EBinOp(lhs, "+", rhs).with_type(e.type)
# yield self.check(lhs.p.body, RUNTIME_POOL)
# yield self.check(rhs.p.body, RUNTIME_POOL)
yield self.check(lhs, RUNTIME_POOL)
yield self.check(rhs, RUNTIME_POOL)
yield self.check(union, RUNTIME_POOL)
# Try instantiating bound expressions
for pool in (STATE_POOL, RUNTIME_POOL):
for (sz1, sz2) in pick_to_sum(2, size - 1):
for e1 in cache.find(pool=pool, size=sz1):
if not is_root(e1):
continue
for v in free_vars(e1):
if pool == RUNTIME_POOL:
e1 = subst(
strip_EStateVar(e1), {
sv.id: EStateVar(sv).with_type(sv.type)
for sv in self.state_vars if sv != v
})
for e2 in cache.find(pool=pool, type=v.type, size=sz2):
yield self.check(subst(e1, {v.id: e2}), pool)
for (sz1, sz2) in pick_to_sum(2, size - 1):
for e in cache.find(pool=RUNTIME_POOL, size=sz1):
if not is_root(e):
continue
for x, pool in map_accelerate(e, self.state_vars, self.binders,
self.args, cache, sz2):
yield self.check(x, pool)
if isinstance(e, EFilter) and not any(v in self.binders
for v in free_vars(e)):
for x, pool in accelerate_filter(e.e, e.p, self.state_vars,
self.binders, self.args,
cache, sz2):
yield self.check(x, pool)
for bag in cache.find_collections(pool=RUNTIME_POOL, size=size - 1):
if not is_root(bag):
continue
for a in self.args:
for v in self.state_vars:
if is_collection(v.type) and v.type == a.type:
v = EStateVar(v).with_type(v.type)
cond = EBinOp(a, BOp.In, v).with_type(BOOL)
yield self.check(
EFilter(bag, mk_lambda(bag.type.t,
lambda _: cond)).with_type(
bag.type), RUNTIME_POOL)
yield self.check(
EFilter(
bag,
mk_lambda(bag.type.t,
lambda _: ENot(cond))).with_type(
bag.type), RUNTIME_POOL)
if isinstance(bag, EFilter):
if any(v not in self.state_vars for v in free_vars(bag.e)):
continue
# separate filter conds
if isinstance(bag.p.body, EBinOp) and bag.p.body.op == BOp.And:
p1 = ELambda(bag.p.arg, bag.p.body.e1)
p2 = ELambda(bag.p.arg, bag.p.body.e2)
f1 = EFilter(bag.e, p1).with_type(bag.type)
f2 = EFilter(bag.e, p2).with_type(bag.type)
f3 = EFilter(f1, p2).with_type(bag.type)
f4 = EFilter(f2, p1).with_type(bag.type)
yield self.check(f1, RUNTIME_POOL)
yield self.check(f2, RUNTIME_POOL)
yield self.check(f3, RUNTIME_POOL)
yield self.check(f4, RUNTIME_POOL)
# construct map lookups
binder = bag.p.arg
inf = infer_map_lookup(bag.p.body, binder,
set(self.state_vars))
if inf:
key_proj, key_lookup, remaining_filter = inf
bag_binder = find_one(
self.binders,
lambda b: b.type == key_proj.type and b != binder)
if bag_binder:
m = strip_EStateVar(
EMakeMap2(
EMap(bag.e,
ELambda(binder, key_proj)).with_type(
type(bag.type)(key_proj.type)),
ELambda(
bag_binder,
EFilter(
bag.e,
ELambda(binder,
EEq(key_proj,
bag_binder))).with_type(
bag.type))).with_type(
TMap(
key_proj.type,
bag.type)))
assert not any(v in self.args for v in free_vars(m))
yield self.check(m, STATE_POOL)
m = EStateVar(m).with_type(m.type)
mg = EMapGet(m, key_lookup).with_type(bag.type)
yield self.check(mg, RUNTIME_POOL)
yield self.check(
EFilter(mg, ELambda(
binder, remaining_filter)).with_type(mg.type),
RUNTIME_POOL)
# for e in cache.find(size=size-1):
# # F(xs +/- ys) ---> F(xs), F(ys)
# for z in break_plus_minus(e):
# if z != e:
# # print("broke {} --> {}".format(pprint(e), pprint(z)))
# yield z
# # try reordering operations
# for (_, e1, f) in enumerate_fragments(e):
# if e1.type == e.type and e1 != e:
# for (_, e2, g) in enumerate_fragments(e1):
# if e2.type == e.type and e2 != e1:
# # e == f(g(e2))
# yield g(f(e2))
yield from self.wrapped.build(cache, size)
def enumerate_core(self, context: Context, size: int, pool: Pool) -> [Exp]:
"""
Arguments:
conext : a Context object describing the vars in scope
size : size to enumerate
pool : pool to enumerate
Yields all expressions of the given size legal in the given context and
pool.
"""
if size < 0:
return
if size == 0:
for (e, p) in LITERALS:
if p == pool:
yield e
for (v, p) in context.vars():
if p == pool:
yield v
for t in all_types(v):
yield construct_value(t)
for (e, ctx, p) in self.hints:
if p == pool and ctx.alpha_equivalent(context):
yield context.adapt(e, ctx)
for t in all_types(e):
yield construct_value(t)
return
yield from self.heuristic_enumeration(context, size, pool)
for e in collections(self.enumerate(context, size - 1, pool)):
yield EEmptyList().with_type(e.type)
if is_numeric(e.type.t):
yield EUnaryOp(UOp.Sum, e).with_type(e.type.t)
for e in self.enumerate(context, size - 1, pool):
yield ESingleton(e).with_type(TBag(e.type))
for e in self.enumerate(context, size - 1, pool):
if isinstance(e.type, TRecord):
for (f, t) in e.type.fields:
yield EGetField(e, f).with_type(t)
for e in self.enumerate(context, size - 1, pool):
if isinstance(e.type, THandle):
yield EGetField(e, "val").with_type(e.type.value_type)
for e in self.enumerate(context, size - 1, pool):
if isinstance(e.type, TTuple):
for n in range(len(e.type.ts)):
yield ETupleGet(e, n).with_type(e.type.ts[n])
for e in of_type(self.enumerate(context, size - 1, pool), BOOL):
yield EUnaryOp(UOp.Not, e).with_type(BOOL)
for e in self.enumerate(context, size - 1, pool):
if is_numeric(e.type):
yield EUnaryOp("-", e).with_type(e.type)
for m in self.enumerate(context, size - 1, pool):
if isinstance(m.type, TMap):
yield EMapKeys(m).with_type(TBag(m.type.k))
for (sz1, sz2) in pick_to_sum(2, size - 1):
for a1 in self.enumerate(context, sz1, pool):
t = a1.type
if not is_numeric(t):
continue
for a2 in of_type(self.enumerate(context, sz2, pool), t):
yield EBinOp(a1, "+", a2).with_type(t)
yield EBinOp(a1, "-", a2).with_type(t)
yield EBinOp(a1, ">", a2).with_type(BOOL)
yield EBinOp(a1, "<", a2).with_type(BOOL)
yield EBinOp(a1, ">=", a2).with_type(BOOL)
yield EBinOp(a1, "<=", a2).with_type(BOOL)
for a1 in collections(self.enumerate(context, sz1, pool)):
for a2 in of_type(self.enumerate(context, sz2, pool), a1.type):
yield EBinOp(a1, "+", a2).with_type(a1.type)
yield EBinOp(a1, "-", a2).with_type(a1.type)
for a2 in of_type(self.enumerate(context, sz2, pool),
a1.type.t):
yield EBinOp(a2, BOp.In, a1).with_type(BOOL)
for a1 in of_type(self.enumerate(context, sz1, pool), BOOL):
for a2 in of_type(self.enumerate(context, sz2, pool), BOOL):
yield EBinOp(a1, BOp.And, a2).with_type(BOOL)
yield EBinOp(a1, BOp.Or, a2).with_type(BOOL)
for a1 in self.enumerate(context, sz1, pool):
if not isinstance(a1.type, TMap):
for a2 in of_type(self.enumerate(context, sz2, pool),
a1.type):
yield EEq(a1, a2)
yield EBinOp(a1, "!=", a2).with_type(BOOL)
for m in self.enumerate(context, sz1, pool):
if isinstance(m.type, TMap):
for k in of_type(self.enumerate(context, sz2, pool),
m.type.k):
yield EMapGet(m, k).with_type(m.type.v)
yield EHasKey(m, k).with_type(BOOL)
for l in self.enumerate(context, sz1, pool):
if not isinstance(l.type, TList):
continue
for i in of_type(self.enumerate(context, sz2, pool), INT):
yield EListGet(l, i).with_type(l.type.t)
for (sz1, sz2, sz3) in pick_to_sum(3, size - 1):
for cond in of_type(self.enumerate(context, sz1, pool), BOOL):
for then_branch in self.enumerate(context, sz2, pool):
for else_branch in of_type(
self.enumerate(context, sz2, pool),
then_branch.type):
yield ECond(cond, then_branch,
else_branch).with_type(then_branch.type)
for l in self.enumerate(context, sz1, pool):
if not isinstance(l.type, TList):
continue
for st in of_type(self.enumerate(context, sz2, pool), INT):
for ed in of_type(self.enumerate(context, sz3, pool), INT):
yield EListSlice(l, st, ed).with_type(l.type)
for bag in collections(self.enumerate(context, size - 1, pool)):
# len of bag
count = EUnaryOp(UOp.Length, bag).with_type(INT)
yield count
# empty?
yield EUnaryOp(UOp.Empty, bag).with_type(BOOL)
# exists?
yield EUnaryOp(UOp.Exists, bag).with_type(BOOL)
# singleton?
yield EEq(count, ONE)
yield EUnaryOp(UOp.The, bag).with_type(bag.type.t)
yield EUnaryOp(UOp.Distinct, bag).with_type(bag.type)
yield EUnaryOp(UOp.AreUnique, bag).with_type(BOOL)
if bag.type.t == BOOL:
yield EUnaryOp(UOp.Any, bag).with_type(BOOL)
yield EUnaryOp(UOp.All, bag).with_type(BOOL)
def build_lambdas(bag, pool, body_size):
v = fresh_var(bag.type.t, omit=set(v for v, p in context.vars()))
inner_context = UnderBinder(context, v=v, bag=bag, bag_pool=pool)
for lam_body in self.enumerate(inner_context, body_size, pool):
yield ELambda(v, lam_body)
# Iteration
for (sz1, sz2) in pick_to_sum(2, size - 1):
for bag in collections(self.enumerate(context, sz1, pool)):
for lam in build_lambdas(bag, pool, sz2):
body_type = lam.body.type
yield EMap(bag, lam).with_type(TBag(body_type))
if body_type == BOOL:
yield EFilter(bag, lam).with_type(bag.type)
if is_numeric(body_type):
yield EArgMin(bag, lam).with_type(bag.type.t)
yield EArgMax(bag, lam).with_type(bag.type.t)
if is_collection(body_type):
yield EFlatMap(bag, lam).with_type(TBag(body_type.t))
# Enable use of a state-pool expression at runtime
if pool == RUNTIME_POOL:
for e in self.enumerate(context, size - 1, STATE_POOL):
yield EStateVar(e).with_type(e.type)
# Create maps
if pool == STATE_POOL:
for (sz1, sz2) in pick_to_sum(2, size - 1):
for bag in collections(self.enumerate(context, sz1,
STATE_POOL)):
if not is_scalar(bag.type.t):
continue
for lam in build_lambdas(bag, STATE_POOL, sz2):
t = TMap(bag.type.t, lam.body.type)
m = EMakeMap2(bag, lam).with_type(t)
yield m
def heuristic_enumeration(self, context: Context, size: int,
pool: Pool) -> [Exp]:
# lambda-instantiation
for sz1, sz2 in pick_to_sum(2, size - 1):
for e in self.enumerate(context, sz1, pool):
for child in e.children():
if isinstance(child, ELambda):
for arg in self.enumerate(context, sz2, pool):
if child.arg.type == arg.type:
yield child.apply_to(arg)
if pool == RUNTIME_POOL:
# is `x` the last of its kind in `xs`?
for sz1, sz2 in pick_to_sum(2, size - 1):
for xs in self.enumerate(context, sz1, STATE_POOL):
if not is_collection(xs.type):
continue
h = histogram(xs)
h = EStateVar(h).with_type(h.type)
for x in of_type(
self.enumerate(context, sz2, RUNTIME_POOL),
xs.type.t):
mg = EMapGet(h, x).with_type(INT)
e = EEq(mg, ONE)
yield e
yield ECond(e,
ESingleton(x).with_type(xs.type),
EEmptyList().with_type(xs.type)).with_type(
xs.type)
yield ECond(e,
EEmptyList().with_type(xs.type),
ESingleton(x).with_type(
xs.type)).with_type(xs.type)
# yield mg
# yield h
# is `x` the last of its kind in `xs` AND is it argmin
# according to some interesting function?
for sz1, sz2 in pick_to_sum(2, size - 1):
for best in self.enumerate(context, sz1, STATE_POOL):
if not (isinstance(best, EArgMin)
or isinstance(best, EArgMax)):
continue
h = histogram(best.e)
h = EStateVar(h).with_type(h.type)
from cozy.structures.heaps import to_heap, EHeapPeek2
heap = to_heap(best)
heap = EStateVar(heap).with_type(heap.type)
for x in of_type(
self.enumerate(context, sz2, RUNTIME_POOL),
best.type):
mg = EMapGet(h, x).with_type(INT)
e = EEq(mg, ONE)
b = EStateVar(best).with_type(best.type)
e = EAll([e, EEq(x, b)])
e._tag = True
yield e
e = ECond(
e,
EHeapPeek2(heap,
EStateVar(ELen(
best.e)).with_type(INT)).with_type(
best.type),
b).with_type(best.type)
e._tag = True
yield e
def test_pick_to_sum_2_3(self):
self.assertEqual(list(pick_to_sum(2, total_size=3)), [(0,3), (1,2), (2,1), (3,0)])
def test_pick_to_sum_2_2(self):
self.assertEqual(list(pick_to_sum(2, total_size=2)), [(0,2), (1,1), (2,0)])
def test_pick_to_sum_1_2(self):
self.assertEqual(list(pick_to_sum(1, total_size=2)), [(2,)])
def test_pick_to_sum_2_3(self):
self.assertEqual(list(pick_to_sum(2, total_size=3)), [(0, 3), (1, 2),
(2, 1), (3, 0)])
def test_pick_to_sum_2_2(self):
self.assertEqual(list(pick_to_sum(2, total_size=2)), [(0, 2), (1, 1),
(2, 0)])
def test_pick_to_sum_1_2(self):
self.assertEqual(list(pick_to_sum(1, total_size=2)), [(2, )])
def _enumerate_core(self, context: Context, size: int,
pool: Pool) -> [Exp]:
"""Build new expressions of the given size.
Arguments:
context : a Context object describing the vars in scope
size : size of expressions to enumerate; each expression in
the output will have this size
pool : pool to enumerate
This function is not cached. Clients should call `enumerate` instead.
This function tries to be a clean description of the Cozy grammar. It
does not concern itself with deduplication (which is handled
efficiently by equivalence class deduplication).
"""
if size < 0:
return
if size == 0:
for e in LITERALS:
yield e
all_interesting_types = OrderedSet(self.hint_types)
for v, _ in context.vars():
all_interesting_types |= all_types(v.type)
for t in all_interesting_types:
l = construct_value(t)
if l not in LITERALS:
yield l
for (v, p) in context.vars():
if p == pool:
yield v
for (e, ctx, p) in self.hints:
if p == pool:
fvs = free_vars(e)
if ctx.alpha_equivalent(context.generalize(fvs)):
yield context.adapt(e, ctx, e_fvs=fvs)
return
if not do_enumerate.value:
return
def build_lambdas(bag, pool, body_size):
v = fresh_var(bag.type.elem_type,
omit=set(v for v, p in context.vars()))
inner_context = UnderBinder(context, v=v, bag=bag, bag_pool=pool)
for lam_body in self.enumerate(inner_context, body_size, pool):
yield ELambda(v, lam_body)
# Load all smaller expressions in this context and pool.
# cache[S] contains expressions of size S in this context and pool.
cache = [list(self.enumerate(context, sz, pool)) for sz in range(size)]
# Enable use of a state-pool expression at runtime
if pool == RUNTIME_POOL:
for e in self.enumerate(context.root(), size - 1, STATE_POOL):
yield EStateVar(e).with_type(e.type)
# Arity-1 expressions
for e in cache[size - 1]:
if is_collection(e.type):
elem_type = e.type.elem_type
# This method of generating EEmptyList() ensures that we visit
# empty collections of all possible types.
yield EEmptyList().with_type(e.type)
if is_numeric(elem_type):
yield EUnaryOp(UOp.Sum, e).with_type(elem_type)
yield EUnaryOp(UOp.Length, e).with_type(INT)
yield EUnaryOp(UOp.Empty, e).with_type(BOOL)
yield EUnaryOp(UOp.Exists, e).with_type(BOOL)
yield EUnaryOp(UOp.The, e).with_type(elem_type)
yield EUnaryOp(UOp.Distinct, e).with_type(e.type)
yield EUnaryOp(UOp.AreUnique, e).with_type(BOOL)
if elem_type == BOOL:
yield EUnaryOp(UOp.Any, e).with_type(BOOL)
yield EUnaryOp(UOp.All, e).with_type(BOOL)
yield ESingleton(e).with_type(TBag(e.type))
if isinstance(e.type, TRecord):
for (f, t) in e.type.fields:
yield EGetField(e, f).with_type(t)
if isinstance(e.type, THandle):
yield EGetField(e, "val").with_type(e.type.value_type)
if isinstance(e.type, TTuple):
for n in range(len(e.type.ts)):
yield ETupleGet(e, n).with_type(e.type.ts[n])
if e.type == BOOL:
yield EUnaryOp(UOp.Not, e).with_type(BOOL)
if is_numeric(e.type):
yield EUnaryOp("-", e).with_type(e.type)
if isinstance(e.type, TMap):
yield EMapKeys(e).with_type(TBag(e.type.k))
# Arity-2 expressions
for (sz1, sz2) in pick_to_sum(2, size - 1):
# sz1 + sz2 = size - 1
for e1 in cache[sz1]:
t = e1.type
if is_numeric(t):
for a2 in of_type(cache[sz2], t):
yield EBinOp(e1, "+", a2).with_type(t)
yield EBinOp(e1, "-", a2).with_type(t)
if is_ordered(t):
for a2 in of_type(cache[sz2], t):
yield EBinOp(e1, ">", a2).with_type(BOOL)
yield EBinOp(e1, "<", a2).with_type(BOOL)
yield EBinOp(e1, ">=", a2).with_type(BOOL)
yield EBinOp(e1, "<=", a2).with_type(BOOL)
if t == BOOL:
for a2 in of_type(cache[sz2], BOOL):
yield EBinOp(e1, BOp.And, a2).with_type(BOOL)
yield EBinOp(e1, BOp.Or, a2).with_type(BOOL)
# Cozy supports the implication operator "=>", but this
# function does not enumerate it because
# - (a => b) is equivalent to ((not a) or b)
# - there isn't an implication operator in any of our
# current target languages, so we would need to
# desugar it to ((not a) or b) anyway.
if not isinstance(t, TMap):
for a2 in of_type(cache[sz2], t):
yield EEq(e1, a2)
yield EBinOp(e1, "!=", a2).with_type(BOOL)
if isinstance(t, TMap):
for k in of_type(cache[sz2], t.k):
yield EMapGet(e1, k).with_type(t.v)
yield EHasKey(e1, k).with_type(BOOL)
if isinstance(t, TList):
for i in of_type(cache[sz2], INT):
yield EListGet(e1, i).with_type(e1.type.elem_type)
if is_collection(t):
elem_type = t.elem_type
for e2 in of_type(cache[sz2], t):
yield EBinOp(e1, "+", e2).with_type(t)
yield EBinOp(e1, "-", e2).with_type(t)
for e2 in of_type(cache[sz2], elem_type):
yield EBinOp(e2, BOp.In, e1).with_type(BOOL)
for f in build_lambdas(e1, pool, sz2):
body_type = f.body.type
yield EMap(e1, f).with_type(TBag(body_type))
if body_type == BOOL:
yield EFilter(e1, f).with_type(t)
if is_numeric(body_type):
yield EArgMin(e1, f).with_type(elem_type)
yield EArgMax(e1, f).with_type(elem_type)
if is_collection(body_type):
yield EFlatMap(e1, f).with_type(
TBag(body_type.elem_type))
if pool == STATE_POOL and is_hashable(elem_type):
yield EMakeMap2(e1, f).with_type(
TMap(elem_type, body_type))
e1_singleton = ESingleton(e1).with_type(TBag(e1.type))
for f in build_lambdas(e1_singleton, pool, sz2):
yield ELet(e1, f).with_type(f.body.type)
# Arity-3 expressions
for (sz1, sz2, sz3) in pick_to_sum(3, size - 1):
# sz1 + sz2 + sz3 = size - 1
for e1 in cache[sz1]:
if e1.type == BOOL:
cond = e1
for then_branch in cache[sz2]:
for else_branch in of_type(cache[sz3],
then_branch.type):
yield ECond(cond, then_branch,
else_branch).with_type(
then_branch.type)
if isinstance(e1.type, TList):
for start in of_type(cache[sz2], INT):
for end in of_type(cache[sz3], INT):
yield EListSlice(e1, start, end).with_type(e1.type)
# It is not necessary to create slice expressions of
# the form a[:i] or a[i:]. Those are desugared
# after parsing to a[0:i] and a[i:len(a)]
# respectively, and Cozy is perfectly capable of
# discovering these expanded forms as well.
for h in all_extension_handlers():
yield from h.enumerate(context, size, pool, self.enumerate,
build_lambdas)