def implement_stmt(self, s : Stm, concretization_functions : { str : Exp }) -> Stm:
op = "<=" if isinstance(s.target.type, TMinHeap) else ">="
f = heap_func(s.target, concretization_functions)
if isinstance(s, SCall):
elem_type = s.target.type.elem_type
target_raw = EVar(s.target.id).with_type(self.rep_type(s.target.type))
if s.func == "add_all":
size = fresh_var(INT, "heap_size")
i = fresh_var(INT, "i")
x = fresh_var(elem_type, "x")
return seq([
SDecl(size.id, s.args[0]),
SEnsureCapacity(target_raw, EBinOp(size, "+", ELen(s.args[1])).with_type(INT)),
SForEach(x, s.args[1], seq([
SAssign(EArrayGet(target_raw, size), x),
SDecl(i.id, size),
SWhile(EAll([
EBinOp(i, ">", ZERO).with_type(BOOL),
ENot(EBinOp(f.apply_to(EArrayGet(target_raw, _parent(i))), op, f.apply_to(EArrayGet(target_raw, i))).with_type(BOOL))]),
seq([
SSwap(EArrayGet(target_raw, _parent(i)), EArrayGet(target_raw, i)),
SAssign(i, _parent(i))])),
SAssign(size, EBinOp(size, "+", ONE).with_type(INT))]))])
elif s.func == "remove_all":
size = fresh_var(INT, "heap_size")
size_minus_one = EBinOp(size, "-", ONE).with_type(INT)
i = fresh_var(INT, "i")
x = fresh_var(elem_type, "x")
label = fresh_name("stop_bubble_down")
child_index = fresh_var(INT, "child_index")
return seq([
SDecl(size.id, s.args[0]),
SForEach(x, s.args[1], seq([
# find the element to remove
SDecl(i.id, EArrayIndexOf(target_raw, x).with_type(INT)),
# swap with last element in heap
SSwap(EArrayGet(target_raw, i), EArrayGet(target_raw, size_minus_one)),
# bubble down
SEscapableBlock(label, SWhile(_has_left_child(i, size_minus_one), seq([
SDecl(child_index.id, _left_child(i)),
SIf(EAll([_has_right_child(i, size_minus_one), ENot(EBinOp(f.apply_to(EArrayGet(target_raw, _left_child(i))), op, f.apply_to(EArrayGet(target_raw, _right_child(i)))))]),
SAssign(child_index, _right_child(i)),
SNoOp()),
SIf(ENot(EBinOp(f.apply_to(EArrayGet(target_raw, i)), op, f.apply_to(EArrayGet(target_raw, child_index)))),
seq([
SSwap(EArrayGet(target_raw, i), EArrayGet(target_raw, child_index)),
SAssign(i, child_index)]),
SEscapeBlock(label))]))),
# dec. size
SAssign(size, size_minus_one)]))])
else:
raise NotImplementedError()
else:
raise NotImplementedError(pprint(s))
def find_one(self, iterable):
v = fresh_var(iterable.type.elem_type, "v")
label = fresh_name("label")
x = fresh_var(iterable.type.elem_type, "x")
decl = SDecl(v, evaluation.construct_value(v.type))
find = SEscapableBlock(label,
SForEach(x, iterable, seq([
SAssign(v, x),
SEscapeBlock(label)])))
self.stms.append(simplify_and_optimize(seq([decl, find])))
return v
def visit_EUnaryOp(self, e):
op = e.op
if op == UOp.Distinct:
return self.visit_iterable(e)
elif op == UOp.The:
return self.find_one(e.e)
elif op == UOp.Sum:
sum_var = fresh_var(e.type, "sum")
loop_var = fresh_var(e.e.type.elem_type, "x")
self.stms.append(simplify_and_optimize(seq([
SDecl(sum_var, ENum(0).with_type(e.type)),
SForEach(loop_var, e.e,
SAssign(sum_var, EBinOp(sum_var, "+", loop_var).with_type(INT)))])))
return sum_var
elif op == UOp.Length:
arg = EVar("x").with_type(e.e.type.elem_type)
return self.visit(EUnaryOp(UOp.Sum, EMap(e.e, ELambda(arg, ONE)).with_type(INT_BAG)).with_type(INT))
elif op == UOp.All:
arg = EVar("x").with_type(e.e.type.elem_type)
return self.visit(EUnaryOp(UOp.Empty, EFilter(e.e, ELambda(arg, ENot(arg))).with_type(INT_BAG)).with_type(INT))
elif op == UOp.Any:
arg = EVar("x").with_type(e.e.type.elem_type)
return self.visit(EUnaryOp(UOp.Exists, EFilter(e.e, ELambda(arg, arg)).with_type(INT_BAG)).with_type(INT))
elif op == UOp.Empty:
iterable = e.e
v = fresh_var(BOOL, "v")
label = fresh_name("label")
x = fresh_var(iterable.type.elem_type, "x")
decl = SDecl(v, ETRUE)
find = SEscapableBlock(label,
SForEach(x, iterable, seq([
SAssign(v, EFALSE),
SEscapeBlock(label)])))
self.stms.append(simplify_and_optimize(seq([decl, find])))
return v
elif op == UOp.Exists:
return self.visit(ENot(EUnaryOp(UOp.Empty, e.e).with_type(BOOL)))
# elif op == UOp.AreUnique:
# s = fresh_var(TSet(e.e.type.elem_type), "unique_elems")
# u = fresh_var(BOOL, "is_unique")
# x = fresh_var(e.e.type.elem_type)
# label = fresh_name("label")
# self.visit(seq([
# SDecl(s, EEmptyList().with_type(s.type)),
# SDecl(u, ETRUE),
# SEscapableBlock(label,
# SForEach(x, e.e,
# SIf(EEscape("{s}.find({x}) != {s}.end()", ("s", "x"), (s, x)).with_type(BOOL),
# seq([SAssign(u, EFALSE), SEscapeBlock(label)]),
# SEscape("{indent}{s}.insert({x});\n", ("s", "x"), (s, x)))))]))
# return u.id
return self.visit_Exp(e)
def visit_EBinOp(self, e):
if e.op in ("+", "-") and is_collection(e.type):
return self.visit_iterable(e)
elif e.op == BOp.In and not isinstance(e.e2.type, TSet):
t = BOOL
res = fresh_var(t, "found")
x = fresh_var(e.e1.type, "x")
label = fresh_name("label")
self.stms.append(simplify_and_optimize(seq([
SDecl(res, EFALSE),
SEscapableBlock(label,
SForEach(x, e.e2, SIf(
EBinOp(x, "==", e.e1).with_type(BOOL),
seq([SAssign(res, ETRUE), SEscapeBlock(label)]),
SNoOp())))])))
return res
return self.visit_Exp(e)
def simplify_and_optimize(s : Stm) -> Stm:
"""Simplify and optimize a statement.
Input:
s - a statement to optimize
Output:
A statement that is functionally equivalent to the input.
This function makes two big transformations:
- "compile" many kinds of expressions (listed below) to simpler forms so
that downstream code generation has less work to do
- avoid creating short-lived intermediate objects (see `stream`)
Expression types eliminated by this procedure:
- EMap, EFilter, EFlatMap
- EArg{Min,Max}
- unary ops:
Distinct,
AreUnique,
Length, Sum, All, Any,
Exists, Empty,
The
- binary ops:
In (where the collection is a Bag or List)
"-" on collections
"+" on collections
- EMakeMap2
- ELet
- EListSlice
- EStm
"""
assert isinstance(s, Stm)
if isinstance(s, SNoOp):
return s
if isinstance(s, SSeq):
# TODO: while the first basic statement in s1 is an SDecl, we should
# apply `re_use` to perhaps eliminate or inline the declaration.
return seq([simplify_and_optimize(s.s1), simplify_and_optimize(s.s2)])
if isinstance(s, SAssign):
setup, e = simplify_and_optimize_expression(s.rhs)
return seq([setup, SAssign(s.lhs, e)])
if isinstance(s, SReturn):
setup, e = simplify_and_optimize_expression(s.e)
return seq([setup, SReturn(e)])
if isinstance(s, SDecl):
setup, e = simplify_and_optimize_expression(s.val)
return seq([setup, SDecl(s.var, e)])
if isinstance(s, SForEach):
return stream(s.iter, s.loop_var, s.body)
if isinstance(s, SEscape):
return s
if isinstance(s, SIf):
setup, test = simplify_and_optimize_expression(s.cond)
if test == ETRUE:
return simplify_and_optimize(s.then_branch)
if test == EFALSE:
return simplify_and_optimize(s.else_branch)
return seq([setup, SIf(test, simplify_and_optimize(s.then_branch), simplify_and_optimize(s.else_branch))])
if isinstance(s, SWhile):
setup, cond = simplify_and_optimize_expression(s.e)
if setup != SNoOp():
# This is a problem because we don't want to duplicate the setup
# condition.
# TODO: introduce an SEscapableBlock/SEscapeBlock to do it right
raise ValueError("oops! setup for condition {} is very long:\n{}".format(pprint(s.e), pprint(setup)))
return SWhile(cond, simplify_and_optimize(s.body))
if isinstance(s, SScoped):
return SScoped(simplify_and_optimize(s.s))
if isinstance(s, SMapUpdate):
# TODO: optimize s.map & s.key
# TODO: s.map must be optimized as an lvalue
mapsetup, map = simplify_and_optimize_lvalue(s.map)
keysetup, key = simplify_and_optimize_expression(s.key)
return seq([
mapsetup,
keysetup,
SMapUpdate(map, key, s.val_var, simplify_and_optimize(s.change))])
if isinstance(s, SMapDel):
mapsetup, map = simplify_and_optimize_lvalue(s.map)
keysetup, key = simplify_and_optimize_expression(s.key)
return seq([
mapsetup,
keysetup,
SMapDel(map, key)])
if isinstance(s, SCall):
setups, args = zip(*(simplify_and_optimize_expression(a) for a in s.args))
return seq(list(setups) + [SCall(s.target, s.func, tuple(args))])
if isinstance(s, SEscapableBlock):
return SEscapableBlock(s.label, simplify_and_optimize(s.body))
if isinstance(s, SEscapeBlock):
return s
if isinstance(s, SArrayAlloc):
setup, cap = simplify_and_optimize_expression(s.capacity)
return seq([setup, SArrayAlloc(s.a, cap)])
if isinstance(s, SArrayReAlloc):
setup, cap = simplify_and_optimize_expression(s.new_capacity)
return seq([setup, SArrayReAlloc(s.a, cap)])
if isinstance(s, SEnsureCapacity):
setup, cap = simplify_and_optimize_expression(s.capacity)
return seq([setup, SEnsureCapacity(s.a, cap)])
if isinstance(s, SSwap):
# TODO: if we want to optimize the operands we will need a special
# procedure that optimizes lvalues while preserving meaning... same
# goes for SAssign case above.
return s
if isinstance(s, SSwitch):
setup, e = simplify_and_optimize_expression(s.e)
new_cases = [(case, simplify_and_optimize(stm)) for (case, stm) in s.cases]
def implement_stmt(self, s: Stm, concretization_functions: {str:
Exp}) -> Stm:
"""Convert a call to a heap function into simpler statements.
This function also requires the `concretization_functions` that
describe the invariants for variables in `e`.
"""
comparison_op = "<=" if isinstance(s.target.type, TMinHeap) else ">="
f = heap_func(s.target, concretization_functions)
if isinstance(s, SCall):
elem_type = s.target.type.elem_type
target_raw = EVar(s.target.id).with_type(
self.rep_type(s.target.type))
if s.func == "add_all":
size = fresh_var(INT, "heap_size")
i = fresh_var(INT, "i")
x = fresh_var(elem_type, "x")
return seq([
SDecl(size.id, s.args[0]),
SEnsureCapacity(
target_raw,
EBinOp(size, "+", ELen(s.args[1])).with_type(INT)),
SForEach(
x, s.args[1],
seq([
SAssign(EArrayGet(target_raw, size), x),
SDecl(i.id, size),
SWhile(
EAll([
EBinOp(i, ">", ZERO).with_type(BOOL),
ENot(
EBinOp(
f.apply_to(
EArrayGet(
target_raw, _parent(i))),
comparison_op,
f.apply_to(EArrayGet(
target_raw,
i))).with_type(BOOL))
]),
seq([
SSwap(EArrayGet(target_raw, _parent(i)),
EArrayGet(target_raw, i)),
SAssign(i, _parent(i))
])),
SAssign(size,
EBinOp(size, "+", ONE).with_type(INT))
]))
])
elif s.func == "remove_all":
size = fresh_var(INT, "heap_size")
size_minus_one = EBinOp(size, "-", ONE).with_type(INT)
i = fresh_var(INT, "i")
x = fresh_var(elem_type, "x")
label = fresh_name("stop_bubble_down")
child_index = fresh_var(INT, "child_index")
return seq([
SDecl(size.id, s.args[0]),
SForEach(
x,
s.args[1],
seq([
# find the element to remove
SDecl(i.id,
EArrayIndexOf(target_raw, x).with_type(INT)),
# swap with last element in heap
SSwap(EArrayGet(target_raw, i),
EArrayGet(target_raw, size_minus_one)),
# bubble down
SEscapableBlock(
label,
SWhile(
_has_left_child(i, size_minus_one),
seq([
SDecl(child_index.id, _left_child(i)),
SIf(
EAll([
_has_right_child(
i, size_minus_one),
ENot(
EBinOp(
f.apply_to(
EArrayGet(
target_raw,
_left_child(
i))),
comparison_op,
f.apply_to(
EArrayGet(
target_raw,
_right_child(
i)))))
]),
SAssign(child_index,
_right_child(i)), SNoOp()),
SIf(
ENot(
EBinOp(
f.apply_to(
EArrayGet(
target_raw, i)),
comparison_op,
f.apply_to(
EArrayGet(
target_raw,
child_index)))),
seq([
SSwap(
EArrayGet(target_raw, i),
EArrayGet(
target_raw,
child_index)),
SAssign(i, child_index)
]), SEscapeBlock(label))
]))),
# dec. size
SAssign(size, size_minus_one)
]))
])
else:
raise ValueError("heaps do not support the function {}".format(
s.func))
else:
raise ValueError(
"the statement {} is not an update to a heap variable".format(
pprint(s)))