def visit_For(self, for_loop):
# Initialize loop variable
self.visit_Assignment(for_loop.init)
# Visit loop statement to get all the written variables
for_written_set = set()
body_ast = AST_C(self.simplify)
# Use the current loop number incremented by one if there is a nested loop inside
body_ast.set_num_loops(self.num_loops + 1)
body_ast.visit(for_loop.stmt)
for_written_set.update(body_ast.written_set)
self.written_set += body_ast.written_set
self.read_set += body_ast.read_set
# The next statement counts as a write
for_written_set.update(body_ast.written_set)
for_written_set.add(for_loop.next.lvalue.name)
increment_id = func.ID(for_loop.next.lvalue.name)
increment_expr = self.expr(for_loop.next.rvalue.__class__, for_loop.next.rvalue)
self.written_set.append(for_loop.next.lvalue.name)
self.written_set.append(for_loop.next.lvalue.name)
outer_id = func.ReturnTuple(for_written_set) if len(for_written_set) > 1 else func.ID(next(iter(for_written_set)))
inner_id = func.ArgsRecList("loop{}".format(self.num_loops), for_written_set)
self.num_loops += 1
# Use a recursive style if statement to replace loop
# If the condition from the loop condition is true, then run the loop body with the increment at the end
# Ignore simplfication for now
cond = self.expr(for_loop.cond.__class__, for_loop.cond)
body_ast.tail_binding.expr2 = func.Binding(increment_id, increment_expr, inner_id)
# Otherwise, return the written variables
if_expr = func.If(cond, body_ast.head_binding, outer_id)
rec_expr = func.RecursiveFunction(inner_id, if_expr, inner_id)
self.__create_binding(outer_id, rec_expr, None)
def visit_If(self, condition):
# Get condition
cond_expr = self.expr(condition.cond.__class__, condition.cond)
self.visit(condition.cond)
iftrue_ast = None
iffalse_ast = None
if_written_set = set()
# When the iftrue or iffalse blocks are not None then visit that branch
# and update the written_set and read_set.
if not condition.iftrue is None:
iftrue_ast = AST_C(self.simplify)
iftrue_ast.visit(condition.iftrue)
self.written_set += iftrue_ast.written_set
self.read_set += iftrue_ast.read_set
if_written_set.update(iftrue_ast.written_set)
if not condition.iffalse is None:
iffalse_ast = AST_C(self.simplify)
iffalse_ast.visit(condition.iffalse)
self.written_set += iffalse_ast.written_set
self.read_set += iffalse_ast.read_set
if_written_set.update(iffalse_ast.written_set)
# Create functional node
lhs = func.ReturnTuple(if_written_set) if len(if_written_set) > 1 else func.ID(next(iter(if_written_set)))
# Set the expression to be the same as the left hand side assignment if the iftrue or iffalse blocks are None.
# Otherwise, set the tail binding's second expression to be the left hand side assignment and set the expression
# as the binding of the associated branch.
if condition.iftrue is None:
iftrue_expr = lhs
else:
iftrue_expr = self.simplify_binding(iftrue_ast.head_binding, iftrue_ast.written_set, iftrue_ast.read_set, if_written_set)
if condition.iffalse is None:
iffalse_expr = lhs
else:
iffalse_expr = self.simplify_binding(iffalse_ast.head_binding, iffalse_ast.written_set, iffalse_ast.read_set, if_written_set)
# Create functional condition node and binding
if_expr = func.If(cond_expr, iftrue_expr, iffalse_expr)
self.__create_binding(lhs, if_expr, None)
def visit_Assignment(self, assignment):
self.visit(assignment.rvalue)
expr1 = self.expr(assignment.rvalue.__class__, assignment.rvalue)
if isinstance(assignment.lvalue, ID):
written_var = func.ID(assignment.lvalue.name)
self.written_set.append(assignment.lvalue.name)
else:
written_var = self.expr(assignment.lvalue.__class__, assignment.lvalue)
# Written Variable is the array name, all the subscripts are read
self.visit_ArrayRef(assignment.lvalue)
expr = assignment.lvalue
while not isinstance(expr, ID):
expr = expr.name
self.written_set.append(expr.name)
self.__create_binding(written_var, expr1, None)
def expr(self, _class, value):
return {
Constant: (lambda orig: func.Constant(orig.value)),
ID: (lambda orig: func.ID(orig.name)),
ArrayRef: (lambda orig: func.ArrayRef(
self.expr(orig.name.__class__, orig.name),
self.expr(orig.subscript.__class__, orig.subscript))),
ExprList: (lambda orig: func.ArgsList([self.expr(x.__class__, x) for x in orig.exprs])),
FuncCall: (lambda orig: func.FuncCall(
self.expr(orig.name.__class__, orig.name),
self.expr(orig.args.__class__, orig.args))),
UnaryOp: (lambda orig: func.UnaryOp(
orig.op,
self.expr(orig.expr.__class__, orig.expr)
)),
BinaryOp: (lambda orig: func.BinaryOp(orig.op,
self.expr(orig.left.__class__ , orig.left),
self.expr(orig.right.__class__, orig.right))),
TernaryOp: (lambda orig: func.If(
self.expr(orig.cond.__class__, orig.cond),
self.expr(orig.iftrue.__class__, orig.iftrue),
self.expr(orig.iffalse.__class__, orig.iffalse)
))
}.get(_class)(value)
def visit_While(self, while_loop):
# Do not need to worry about incrementation and initialization in while loop. Assume they're there and loop can terminate.
# Visit loop statement to get all the written variables
while_written_set = set()
body_ast = AST_C(self.simplify)
body_ast.set_num_loops(self.num_loops + 1)
body_ast.visit(while_loop.stmt)
while_written_set.update(body_ast.written_set)
self.written_set += body_ast.written_set
self.read_set += body_ast.read_set
outer_id = func.ReturnTuple(while_written_set) if len(while_written_set) > 1 else func.ID(next(iter(while_written_set)))
inner_id = func.ArgsRecList("loop{}".format(self.num_loops), while_written_set)
self.num_loops += 1
# Use a recursive style if statement to replace loop
# If the condition from the loop condition is true, then run the loop body with the increment at the end
# Ignore simplfication for now
cond = self.expr(while_loop.cond.__class__, while_loop.cond)
body_ast.tail_binding.expr2 = inner_id
if_expr = func.If(cond, body_ast.head_binding, outer_id)
rec_expr = func.RecursiveFunction(inner_id, if_expr, inner_id)
self.__create_binding(outer_id, rec_expr, None)