229 lines
10 KiB
Python
229 lines
10 KiB
Python
from antlr4.tree.Tree import TerminalNode
|
|
|
|
from .antlr_router import get_antlr_modules
|
|
|
|
_, P, _ = get_antlr_modules()
|
|
|
|
# builtins that print, touch the stack or change their argument in place; calls to them are never removed
|
|
_EFFECT_TOKENS = {
|
|
P.PRNT, P.PRINT_SHAPE, P.PUSH, P.POP, P.CLEAR, P.HAS, P.GET,
|
|
P.ADD_KEY, P.REMOVE_KEY, P.ENABLE_TRACKING, P.DISABLE_TRACKING, P.RESET_TRACKING,
|
|
}
|
|
|
|
|
|
def remove_dead_code(tree):
|
|
"""Removes statements whose result is never read and that have no side effects. Changes and returns tree."""
|
|
_DeadCode(tree).run()
|
|
return tree
|
|
|
|
|
|
class _DeadCode:
|
|
def __init__(self, tree):
|
|
self.tree = tree
|
|
self.defs = {}
|
|
self.lambdas = []
|
|
self.calls = []
|
|
assignments = []
|
|
# which lambdas a variable can hold, and variables that may hold one we cannot follow
|
|
self.lam = {}
|
|
self.unknown = set()
|
|
todo = [tree]
|
|
while todo:
|
|
n = todo.pop()
|
|
if isinstance(n, TerminalNode):
|
|
continue
|
|
if isinstance(n, P.VarDefContext):
|
|
if n.EQUEALS() and not n.LBRACKET():
|
|
assignments.append((n.VARIABLE().getText(), n.expr(0)))
|
|
elif isinstance(n, P.ForStmtContext):
|
|
self.unknown.add(n.VARIABLE().getText())
|
|
elif isinstance(n, P.FunctionDefContext):
|
|
self.defs.setdefault(n.VARIABLE().getText(), []).append((_params(n), n.block() or n.expr()))
|
|
elif isinstance(n, P.LambdaExpContext):
|
|
self.lambdas.append(n)
|
|
elif isinstance(n, P.CallExpContext):
|
|
self.calls.append(n)
|
|
todo.extend(n.getChildren())
|
|
|
|
# User functions see the caller's variables, so a call reads everything its body reads.
|
|
# Lambdas only see the variables captured when they were created, so calling one reads nothing extra.
|
|
self.fn_reads = {name: set() for name in self.defs}
|
|
self.fn_effect = {name: False for name in self.defs}
|
|
self.lambda_effect = {lam: False for lam in self.lambdas}
|
|
changed = True
|
|
while changed:
|
|
changed = False
|
|
for name, expr in assignments:
|
|
changed |= self.flow(name, self.value(expr))
|
|
for call in self.calls:
|
|
lams, fns, _ = self.resolve(call.VARIABLE().getText())
|
|
args = call.exprList().expr() if call.exprList() else []
|
|
targets = [_params(lam) for lam in lams] + [params for fn in fns for params, _ in self.defs[fn]]
|
|
for params in targets:
|
|
for param, arg in zip(params, args):
|
|
changed |= self.flow(param, self.value(arg))
|
|
for lam in self.lambdas:
|
|
if not self.lambda_effect[lam] and self.scan(lam.block() or lam.expr())[1]:
|
|
self.lambda_effect[lam] = changed = True
|
|
for name, defs in self.defs.items():
|
|
reads, effect = set(), False
|
|
for params, body in defs:
|
|
r, e = self.scan(body)
|
|
reads |= r - set(params)
|
|
effect |= e
|
|
if reads != self.fn_reads[name] or effect != self.fn_effect[name]:
|
|
self.fn_reads[name], self.fn_effect[name] = reads, effect
|
|
changed = True
|
|
|
|
def flow(self, name, value):
|
|
lams, unknown = value
|
|
changed = not lams <= self.lam.get(name, set()) or (unknown and name not in self.unknown)
|
|
self.lam.setdefault(name, set()).update(lams)
|
|
if unknown:
|
|
self.unknown.add(name)
|
|
return changed
|
|
|
|
def value(self, node):
|
|
"""Lambdas an expression can evaluate to, and whether it may be one we cannot follow."""
|
|
while True:
|
|
if isinstance(node, P.ParenExpContext):
|
|
node = node.expr()
|
|
elif node.getChildCount() == 1 and not isinstance(node.getChild(0), TerminalNode):
|
|
node = node.getChild(0)
|
|
else:
|
|
break
|
|
if isinstance(node, P.LambdaExpContext):
|
|
return {node}, False
|
|
if isinstance(node, P.VariableExpContext):
|
|
name = node.getText()
|
|
return set(self.lam.get(name, ())), name in self.unknown
|
|
if isinstance(node, P.TernaryExpContext):
|
|
(a, ua), (b, ub) = self.value(node.expr(0)), self.value(node.expr(1))
|
|
return a | b, ua or ub
|
|
todo = [node]
|
|
while todo:
|
|
n = todo.pop()
|
|
if isinstance(n, TerminalNode):
|
|
continue
|
|
if isinstance(n, (P.LambdaExpContext, P.CallExpContext)):
|
|
return set(), True
|
|
if isinstance(n, P.VariableExpContext) and (self.lam.get(n.getText()) or n.getText() in self.unknown):
|
|
return set(), True
|
|
todo.extend(n.getChildren())
|
|
return set(), False
|
|
|
|
def resolve(self, name):
|
|
"""Lambdas and user functions a call by this name can reach, and whether it may reach something else."""
|
|
lams = self.lam.get(name, set())
|
|
unknown = name in self.unknown
|
|
fns = [name] if name in self.defs else []
|
|
return lams, fns, unknown or not (lams or fns)
|
|
|
|
def run(self):
|
|
live, _ = self.scan(self.tree.expr())
|
|
self.sweep(self.tree, live, None, apply=True, keep_last=False, skip=self.tree.expr())
|
|
|
|
def scan(self, node):
|
|
"""Variables an expression reads, and whether evaluating it has side effects."""
|
|
reads, effect = set(), False
|
|
todo = [node]
|
|
while todo:
|
|
n = todo.pop()
|
|
if isinstance(n, TerminalNode):
|
|
t = n.getSymbol().type
|
|
if t == P.VARIABLE:
|
|
reads.add(n.getText())
|
|
elif t in _EFFECT_TOKENS:
|
|
effect = True
|
|
continue
|
|
if isinstance(n, P.LambdaExpContext):
|
|
# creating a lambda only captures variables; its body runs when it is called
|
|
reads |= self.scan(n.block() or n.expr())[0] - set(_params(n))
|
|
continue
|
|
if isinstance(n, P.CallExpContext):
|
|
lams, fns, unknown = self.resolve(n.VARIABLE().getText())
|
|
effect |= unknown or any(self.lambda_effect[lam] for lam in lams) or any(self.fn_effect[fn] for fn in fns)
|
|
for fn in fns:
|
|
reads |= self.fn_reads[fn]
|
|
todo.extend(n.getChildren())
|
|
return reads, effect
|
|
|
|
def sweep(self, parent, live, loop_live, apply, keep_last, skip=None):
|
|
"""Walks a statement list backwards, removing dead statements when apply is set; returns live-before."""
|
|
stmts = [c for c in parent.children if not isinstance(c, TerminalNode) and c is not skip]
|
|
for i in range(len(stmts) - 1, -1, -1):
|
|
before, removable = self.stmt(stmts[i], live, loop_live, apply)
|
|
if removable and not (keep_last and i == len(stmts) - 1):
|
|
if apply:
|
|
parent.children.remove(stmts[i])
|
|
else:
|
|
live = before
|
|
return live
|
|
|
|
def stmt(self, node, live, loop_live, apply):
|
|
"""Returns (live-before if kept, whether the statement can be removed)."""
|
|
if isinstance(node, P.VarDefStmtContext):
|
|
node = node.varDef()
|
|
if isinstance(node, P.VarDefContext):
|
|
name = node.VARIABLE().getText()
|
|
reads, effect = set(), False
|
|
for e in node.expr():
|
|
r, eff = self.scan(e)
|
|
reads |= r
|
|
effect |= eff
|
|
if node.EQUEALS() and not node.LBRACKET():
|
|
before = (live - {name}) | reads
|
|
else:
|
|
before = live | {name} | reads
|
|
return before, name not in live and not effect
|
|
if isinstance(node, P.ExprStatementContext):
|
|
reads, effect = self.scan(node.expr())
|
|
return live | reads, not effect
|
|
if isinstance(node, P.IfStatementContext):
|
|
node = node.ifStmt()
|
|
reads, _ = self.scan(node.expr())
|
|
before = reads | self.stmt_kept(node.stmt(0), live, loop_live, apply)
|
|
if node.stmt(1):
|
|
return before | self.stmt_kept(node.stmt(1), live, loop_live, apply), False
|
|
return before | live, False
|
|
if isinstance(node, P.WhileStatementContext):
|
|
node = node.whileStmt()
|
|
reads, _ = self.scan(node.expr())
|
|
head = self.loop_head(node.stmt(), live | reads, None, apply)
|
|
return head, False
|
|
if isinstance(node, P.ForStatementContext):
|
|
node = node.forStmt()
|
|
reads, _ = self.scan(node.expr())
|
|
head = self.loop_head(node.stmt(), live, node.VARIABLE().getText(), apply)
|
|
return head | reads, False
|
|
if isinstance(node, P.BlockStatementContext):
|
|
return self.sweep(node.block(), live, loop_live, apply, keep_last=True), False
|
|
if isinstance(node, (P.BreakStatementContext, P.ContinueStatementContext)):
|
|
return set(loop_live), False
|
|
if isinstance(node, P.ReturnStatementContext):
|
|
expr = node.returnStmt().expr()
|
|
return (self.scan(expr)[0] if expr else set()), False
|
|
# function definitions and anything unexpected stay, and are assumed to read everything still live
|
|
return live, False
|
|
|
|
def stmt_kept(self, node, live, loop_live, apply):
|
|
return self.stmt(node, live, loop_live, apply)[0]
|
|
|
|
def loop_head(self, body, exit_live, loop_var, apply):
|
|
"""Fixed point of the variables live at the top of a loop; only removes once it is stable."""
|
|
head = set(exit_live)
|
|
while True:
|
|
body_live = self.stmt_kept(body, head, head, apply=False)
|
|
if loop_var is not None:
|
|
body_live = body_live - {loop_var}
|
|
new_head = exit_live | body_live
|
|
if new_head == head:
|
|
break
|
|
head = new_head
|
|
self.stmt_kept(body, head, head, apply)
|
|
return head
|
|
|
|
|
|
def _params(node):
|
|
return [v.getText() for v in node.paramList().VARIABLE()] if node.paramList() else []
|