818 lines
26 KiB
Python
818 lines
26 KiB
Python
|
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
||
|
#
|
||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
# you may not use this file except in compliance with the License.
|
||
|
# You may obtain a copy of the License at
|
||
|
#
|
||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
#
|
||
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
# See the License for the specific language governing permissions and
|
||
|
# limitations under the License.
|
||
|
# ==============================================================================
|
||
|
"""Control flow graph (CFG) structure for Python AST representation.
|
||
|
|
||
|
The CFG is a digraph with edges representing valid control flow. Each
|
||
|
node is associated with exactly one AST node, but not all AST nodes may have
|
||
|
a corresponding CFG counterpart.
|
||
|
|
||
|
Once built, the CFG itself is immutable, but the values it holds need not be;
|
||
|
they are usually annotated with information extracted by walking the graph.
|
||
|
"""
|
||
|
|
||
|
from __future__ import absolute_import
|
||
|
from __future__ import division
|
||
|
from __future__ import print_function
|
||
|
|
||
|
import collections
|
||
|
from enum import Enum
|
||
|
|
||
|
# pylint:disable=g-bad-import-order
|
||
|
import gast
|
||
|
# pylint:enable=g-bad-import-order
|
||
|
|
||
|
from tensorflow.contrib.autograph.pyct import compiler
|
||
|
|
||
|
|
||
|
class Node(object):
|
||
|
"""A node in the CFG.
|
||
|
|
||
|
Although new instances of this class are mutable, the objects that a user
|
||
|
finds in the CFG are typically not.
|
||
|
|
||
|
The nodes represent edges in the CFG graph, and maintain pointers to allow
|
||
|
efficient walking in both forward and reverse order. The following property
|
||
|
holds for all nodes: "child in node.next" iff "node in child.prev".
|
||
|
|
||
|
Attributes:
|
||
|
next: FrozenSet[Node, ...], the nodes that follow this node, in control
|
||
|
flow order
|
||
|
prev: FrozenSet[Node, ...], the nodes that precede this node, in reverse
|
||
|
control flow order
|
||
|
ast_node: ast.AST, the AST node corresponding to this CFG node
|
||
|
"""
|
||
|
|
||
|
def __init__(self, next_, prev, ast_node):
|
||
|
self.next = next_
|
||
|
self.prev = prev
|
||
|
self.ast_node = ast_node
|
||
|
|
||
|
def freeze(self):
|
||
|
self.next = frozenset(self.next)
|
||
|
self.prev = frozenset(self.prev)
|
||
|
|
||
|
def __repr__(self):
|
||
|
if isinstance(self.ast_node, gast.FunctionDef):
|
||
|
return 'def %s' % self.ast_node.name
|
||
|
elif isinstance(self.ast_node, gast.withitem):
|
||
|
source, _ = compiler.ast_to_source(self.ast_node.context_expr)
|
||
|
return source.strip()
|
||
|
source, _ = compiler.ast_to_source(self.ast_node)
|
||
|
return source.strip()
|
||
|
|
||
|
|
||
|
class Graph(
|
||
|
collections.namedtuple(
|
||
|
'Graph',
|
||
|
['entry', 'exit', 'error', 'index', 'stmt_prev', 'stmt_next'])):
|
||
|
"""A Control Flow Graph.
|
||
|
|
||
|
The CFG maintains an index to allow looking up a CFG node by the AST node to
|
||
|
which it is associated. The index can also be enumerated in top-down, depth
|
||
|
first order.
|
||
|
|
||
|
Walking the graph in forward or reverse order is supported by double
|
||
|
parent-child links.
|
||
|
|
||
|
Note: the error nodes are not wired to their corresponding finally guards,
|
||
|
because these are shared, and wiring them would create a reverse path from
|
||
|
normal control flow into the error nodes, which we want to avoid.
|
||
|
|
||
|
The graph also maintains edges corresponding to higher level statements
|
||
|
like for-else loops. A node is considered successor of a statement if there
|
||
|
is an edge from a node that is lexically a child of that statement to a node
|
||
|
that is not. Statement predecessors are analogously defined.
|
||
|
|
||
|
Attributes:
|
||
|
entry: Node, the entry node
|
||
|
exit: FrozenSet[Node, ...], the exit nodes
|
||
|
error: FrozenSet[Node, ...], nodes that exit due to an explicitly raised
|
||
|
error (errors propagated from function calls are not accounted)
|
||
|
index: Dict[ast.Node, Node], mapping AST nodes to the respective CFG
|
||
|
node
|
||
|
stmt_prev: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST
|
||
|
nodes to their predecessor CFG nodes
|
||
|
stmt_next: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST
|
||
|
nodes to their successor CFG nodes
|
||
|
"""
|
||
|
|
||
|
def __repr__(self):
|
||
|
result = 'digraph CFG {\n'
|
||
|
for node in self.index.values():
|
||
|
result += ' %s [label="%s"];\n' % (id(node), node)
|
||
|
for node in self.index.values():
|
||
|
for next_ in node.next:
|
||
|
result += ' %s -> %s;\n' % (id(node), id(next_))
|
||
|
result += '}'
|
||
|
return result
|
||
|
|
||
|
|
||
|
class _WalkMode(Enum):
|
||
|
FORWARD = 1
|
||
|
REVERSE = 2
|
||
|
|
||
|
|
||
|
# TODO(mdan): Rename to DataFlowAnalyzer.
|
||
|
# TODO(mdan): Consider specializations that use gen/kill/transfer abstractions.
|
||
|
class GraphVisitor(object):
|
||
|
"""Base class for a CFG visitors.
|
||
|
|
||
|
This implementation is not thread safe.
|
||
|
|
||
|
The visitor has some facilities to simplify dataflow analyses. In particular,
|
||
|
it allows revisiting the nodes at the decision of the subclass. This can be
|
||
|
used to visit the graph until the state reaches a fixed point.
|
||
|
|
||
|
For more details on dataflow analysis, see
|
||
|
https://www.seas.harvard.edu/courses/cs252/2011sp/slides/Lec02-Dataflow.pdf
|
||
|
|
||
|
Note: the literature generally suggests visiting successor nodes only when the
|
||
|
state of the current node changed, regardless of whether that successor has
|
||
|
ever been visited. This implementation visits every successor at least once.
|
||
|
|
||
|
Attributes:
|
||
|
graph: Graph
|
||
|
in_: Dict[Node, Any], stores node-keyed state during a visit
|
||
|
out: Dict[Node, Any], stores node-keyed state during a visit
|
||
|
"""
|
||
|
|
||
|
def __init__(self, graph):
|
||
|
self.graph = graph
|
||
|
self.reset()
|
||
|
|
||
|
def init_state(self, node):
|
||
|
"""State initialization function. Optional to overload.
|
||
|
|
||
|
An in/out state slot will be created for each node in the graph. Subclasses
|
||
|
must overload this to control what that is initialized to.
|
||
|
|
||
|
Args:
|
||
|
node: Node
|
||
|
"""
|
||
|
raise NotImplementedError('Subclasses must implement this.')
|
||
|
|
||
|
# TODO(mdan): Rename to flow?
|
||
|
def visit_node(self, node):
|
||
|
"""Visitor function.
|
||
|
|
||
|
Args:
|
||
|
node: Node
|
||
|
Returns:
|
||
|
bool, whether the node should be revisited; subclasses can visit every
|
||
|
reachable node exactly once by always returning False
|
||
|
"""
|
||
|
raise NotImplementedError('Subclasses must implement this.')
|
||
|
|
||
|
def reset(self):
|
||
|
self.in_ = {
|
||
|
node: self.init_state(node) for node in self.graph.index.values()
|
||
|
}
|
||
|
self.out = {
|
||
|
node: self.init_state(node) for node in self.graph.index.values()
|
||
|
}
|
||
|
|
||
|
def _visit_internal(self, mode):
|
||
|
"""Visits the CFG, depth-first."""
|
||
|
assert mode in (_WalkMode.FORWARD, _WalkMode.REVERSE)
|
||
|
if mode == _WalkMode.FORWARD:
|
||
|
open_ = [self.graph.entry]
|
||
|
elif mode == _WalkMode.REVERSE:
|
||
|
open_ = list(self.graph.exit)
|
||
|
closed = set()
|
||
|
|
||
|
while open_:
|
||
|
node = open_.pop(0)
|
||
|
closed.add(node)
|
||
|
|
||
|
should_revisit = self.visit_node(node)
|
||
|
|
||
|
if mode == _WalkMode.FORWARD:
|
||
|
children = node.next
|
||
|
elif mode == _WalkMode.REVERSE:
|
||
|
children = node.prev
|
||
|
|
||
|
for next_ in children:
|
||
|
if should_revisit or next_ not in closed:
|
||
|
open_.append(next_)
|
||
|
|
||
|
def visit_forward(self):
|
||
|
self._visit_internal(_WalkMode.FORWARD)
|
||
|
|
||
|
def visit_reverse(self):
|
||
|
self._visit_internal(_WalkMode.REVERSE)
|
||
|
|
||
|
|
||
|
class GraphBuilder(object):
|
||
|
"""Builder that constructs a CFG from a given AST.
|
||
|
|
||
|
This GraphBuilder facilitates constructing the DAG that forms the CFG when
|
||
|
nodes
|
||
|
are supplied in lexical order (i.e., top-down, depth first). Under these
|
||
|
conditions, it supports building patterns found in typical structured
|
||
|
programs.
|
||
|
|
||
|
This builder ignores the flow generated by exceptions, which are assumed to
|
||
|
always be catastrophic and present purely for diagnostic purposes (e.g. to
|
||
|
print debug information). Statements like raise and try/catch sections are
|
||
|
allowed and will generate control flow edges, but ordinaty statements are
|
||
|
assumed not to raise exceptions.
|
||
|
|
||
|
Finally sections are also correctly interleaved between break/continue/return
|
||
|
nodes and their subsequent statements.
|
||
|
|
||
|
Important concepts:
|
||
|
* nodes - nodes refer refer to CFG nodes; AST nodes are qualified explicitly
|
||
|
* leaf set - since the graph is constructed gradually, a leaf set maintains
|
||
|
the CFG nodes that will precede the node that the builder expects to
|
||
|
receive next; when an ordinary node is added, it is connected to the
|
||
|
existing leaves and it in turn becomes the new leaf
|
||
|
* jump nodes - nodes that should generate edges other than what
|
||
|
ordinary nodes would; these correspond to break, continue and return
|
||
|
statements
|
||
|
* sections - logical delimiters for subgraphs that require special
|
||
|
edges; there are various types of nodes, each admitting various
|
||
|
types of jump nodes; sections are identified by their corresponding AST
|
||
|
node
|
||
|
"""
|
||
|
|
||
|
# TODO(mdan): Perhaps detail this in a markdown doc.
|
||
|
# TODO(mdan): Add exception support.
|
||
|
|
||
|
def __init__(self, parent_ast_node):
|
||
|
self.reset()
|
||
|
self.parent = parent_ast_node
|
||
|
|
||
|
def reset(self):
|
||
|
"""Resets the state of this factory."""
|
||
|
self.head = None
|
||
|
self.errors = set()
|
||
|
self.node_index = collections.OrderedDict()
|
||
|
|
||
|
# TODO(mdan): Too many primitives. Use classes.
|
||
|
self.leaves = set()
|
||
|
|
||
|
# Note: This mechanism requires that nodes are added in lexical order (top
|
||
|
# to bottom, depth first).
|
||
|
self.active_stmts = set()
|
||
|
self.owners = {} # type: Set[any]
|
||
|
self.forward_edges = set() # type: Tuple[Node, Node] # (from, to)
|
||
|
|
||
|
self.finally_sections = {}
|
||
|
# Dict values represent (entry, exits)
|
||
|
self.finally_section_subgraphs = {
|
||
|
} # type: Dict[ast.AST, Tuple[Node, Set[Node]]]
|
||
|
# Whether the guard section can be reached from the statement that precedes
|
||
|
# it.
|
||
|
self.finally_section_has_direct_flow = {}
|
||
|
# Finally sections that await their first node.
|
||
|
self.pending_finally_sections = set()
|
||
|
|
||
|
# Exit jumps keyed by the section they affect.
|
||
|
self.exits = {}
|
||
|
|
||
|
# The entry of loop sections, keyed by the section.
|
||
|
self.section_entry = {}
|
||
|
# Continue jumps keyed by the section they affect.
|
||
|
self.continues = {}
|
||
|
|
||
|
# The entry of conditional sections, keyed by the section.
|
||
|
self.cond_entry = {}
|
||
|
# Lists of leaf nodes corresponding to each branch in the section.
|
||
|
self.cond_leaves = {}
|
||
|
|
||
|
def _connect_nodes(self, first, second):
|
||
|
"""Connects nodes to signify that control flows from first to second.
|
||
|
|
||
|
Args:
|
||
|
first: Union[Set[Node, ...], Node]
|
||
|
second: Node
|
||
|
"""
|
||
|
if isinstance(first, Node):
|
||
|
first.next.add(second)
|
||
|
second.prev.add(first)
|
||
|
self.forward_edges.add((first, second))
|
||
|
else:
|
||
|
for node in first:
|
||
|
self._connect_nodes(node, second)
|
||
|
|
||
|
def _add_new_node(self, ast_node):
|
||
|
"""Grows the graph by adding a CFG node following the current leaves."""
|
||
|
if ast_node is self.node_index:
|
||
|
raise ValueError('%s added twice' % ast_node)
|
||
|
node = Node(next_=set(), prev=set(), ast_node=ast_node)
|
||
|
self.node_index[ast_node] = node
|
||
|
self.owners[node] = frozenset(self.active_stmts)
|
||
|
|
||
|
if self.head is None:
|
||
|
self.head = node
|
||
|
|
||
|
for leaf in self.leaves:
|
||
|
self._connect_nodes(leaf, node)
|
||
|
|
||
|
# If any finally section awaits its first node, populate it.
|
||
|
for section_id in self.pending_finally_sections:
|
||
|
self.finally_section_subgraphs[section_id][0] = node
|
||
|
self.pending_finally_sections = set()
|
||
|
|
||
|
return node
|
||
|
|
||
|
def begin_statement(self, stmt):
|
||
|
"""Marks the beginning of a statement.
|
||
|
|
||
|
Args:
|
||
|
stmt: Hashable, a key by which the statement can be identified in
|
||
|
the CFG's stmt_prev and stmt_next attributes
|
||
|
"""
|
||
|
self.active_stmts.add(stmt)
|
||
|
|
||
|
def end_statement(self, stmt):
|
||
|
"""Marks the end of a statement.
|
||
|
|
||
|
Args:
|
||
|
stmt: Hashable, a key by which the statement can be identified in
|
||
|
the CFG's stmt_prev and stmt_next attributes; must match a key
|
||
|
previously passed to begin_statement.
|
||
|
"""
|
||
|
self.active_stmts.remove(stmt)
|
||
|
|
||
|
def add_ordinary_node(self, ast_node):
|
||
|
"""Grows the graph by adding an ordinary CFG node.
|
||
|
|
||
|
Ordinary nodes are followed by the next node, in lexical order, that is,
|
||
|
they become the new leaf set.
|
||
|
|
||
|
Args:
|
||
|
ast_node: ast.AST
|
||
|
Returns:
|
||
|
Node
|
||
|
"""
|
||
|
node = self._add_new_node(ast_node)
|
||
|
self.leaves = set((node,))
|
||
|
return node
|
||
|
|
||
|
def _add_jump_node(self, ast_node, guards):
|
||
|
"""Grows the graph by adding a jump node.
|
||
|
|
||
|
Jump nodes are added to the current leaf set, and the leaf set becomes
|
||
|
empty. If the jump node is the last in a cond section, then it may be added
|
||
|
back to the leaf set by a separate mechanism.
|
||
|
|
||
|
Args:
|
||
|
ast_node: ast.AST
|
||
|
guards: Tuple[ast.AST, ...], the finally sections active for this node
|
||
|
Returns:
|
||
|
Node
|
||
|
"""
|
||
|
node = self._add_new_node(ast_node)
|
||
|
self.leaves = set()
|
||
|
# The guards themselves may not yet be complete, and will be wired later.
|
||
|
self.finally_sections[node] = guards
|
||
|
return node
|
||
|
|
||
|
def _connect_jump_to_finally_sections(self, node):
|
||
|
"""Connects a jump node to the finally sections protecting it."""
|
||
|
cursor = set((node,))
|
||
|
for guard_section_id in self.finally_sections[node]:
|
||
|
guard_begin, guard_ends = self.finally_section_subgraphs[guard_section_id]
|
||
|
self._connect_nodes(cursor, guard_begin)
|
||
|
cursor = guard_ends
|
||
|
del self.finally_sections[node]
|
||
|
# TODO(mdan): Should garbage-collect finally_section_subgraphs.
|
||
|
return cursor
|
||
|
|
||
|
def add_exit_node(self, ast_node, section_id, guards):
|
||
|
"""Grows the graph by adding an exit node.
|
||
|
|
||
|
This node becomes an exit for the current section.
|
||
|
|
||
|
Args:
|
||
|
ast_node: ast.AST
|
||
|
section_id: Hashable, the node for which ast_node should be considered
|
||
|
to be an exit node
|
||
|
guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
|
||
|
"""
|
||
|
node = self._add_jump_node(ast_node, guards)
|
||
|
self.exits[section_id].add(node)
|
||
|
|
||
|
def add_continue_node(self, ast_node, section_id, guards):
|
||
|
"""Grows the graph by adding a reentry node.
|
||
|
|
||
|
This node causes control flow to go back to the loop section's entry.
|
||
|
|
||
|
Args:
|
||
|
ast_node: ast.AST
|
||
|
section_id: Hashable, the node for which ast_node should be considered
|
||
|
to be an exit node
|
||
|
guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
|
||
|
"""
|
||
|
node = self._add_jump_node(ast_node, guards)
|
||
|
self.continues[section_id].add(node)
|
||
|
|
||
|
def add_error_node(self, ast_node, guards):
|
||
|
"""Grows the graph by adding an error node.
|
||
|
|
||
|
This node becomes an exit for the entire graph.
|
||
|
|
||
|
Args:
|
||
|
ast_node: ast.AST
|
||
|
guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
|
||
|
"""
|
||
|
node = self._add_jump_node(ast_node, guards)
|
||
|
self.errors.add(node)
|
||
|
self.leaves = set()
|
||
|
|
||
|
def enter_section(self, section_id):
|
||
|
"""Enters a regular section.
|
||
|
|
||
|
Regular sections admit exit jumps, which end the section.
|
||
|
|
||
|
Args:
|
||
|
section_id: Hashable, the same node that will be used in calls to the
|
||
|
ast_node arg passed to add_exit_node
|
||
|
"""
|
||
|
assert section_id not in self.exits
|
||
|
self.exits[section_id] = set()
|
||
|
|
||
|
def exit_section(self, section_id):
|
||
|
"""Exits a regular section."""
|
||
|
|
||
|
# Exits are jump nodes, which may be protected.
|
||
|
for exit_ in self.exits[section_id]:
|
||
|
self.leaves |= self._connect_jump_to_finally_sections(exit_)
|
||
|
|
||
|
del self.exits[section_id]
|
||
|
|
||
|
def enter_loop_section(self, section_id, entry_node):
|
||
|
"""Enters a loop section.
|
||
|
|
||
|
Loop sections define an entry node. The end of the section always flows back
|
||
|
to the entry node. These admit continue jump nodes which also flow to the
|
||
|
entry node.
|
||
|
|
||
|
Args:
|
||
|
section_id: Hashable, the same node that will be used in calls to the
|
||
|
ast_node arg passed to add_continue_node
|
||
|
entry_node: ast.AST, the entry node into the loop (e.g. the test node
|
||
|
for while loops)
|
||
|
"""
|
||
|
assert section_id not in self.section_entry
|
||
|
assert section_id not in self.continues
|
||
|
self.continues[section_id] = set()
|
||
|
node = self.add_ordinary_node(entry_node)
|
||
|
self.section_entry[section_id] = node
|
||
|
|
||
|
def exit_loop_section(self, section_id):
|
||
|
"""Exits a loop section."""
|
||
|
self._connect_nodes(self.leaves, self.section_entry[section_id])
|
||
|
|
||
|
# continues are jump nodes, which may be protected.
|
||
|
for reentry in self.continues[section_id]:
|
||
|
guard_ends = self._connect_jump_to_finally_sections(reentry)
|
||
|
self._connect_nodes(guard_ends, self.section_entry[section_id])
|
||
|
|
||
|
# Loop nodes always loop back.
|
||
|
self.leaves = set((self.section_entry[section_id],))
|
||
|
|
||
|
del self.continues[section_id]
|
||
|
del self.section_entry[section_id]
|
||
|
|
||
|
def enter_cond_section(self, section_id):
|
||
|
"""Enters a conditional section.
|
||
|
|
||
|
Conditional sections define an entry node, and one or more branches.
|
||
|
|
||
|
Args:
|
||
|
section_id: Hashable, the same node that will be used in calls to the
|
||
|
section_id arg passed to new_cond_branch
|
||
|
"""
|
||
|
|
||
|
assert section_id not in self.cond_entry
|
||
|
assert section_id not in self.cond_leaves
|
||
|
self.cond_leaves[section_id] = []
|
||
|
|
||
|
def new_cond_branch(self, section_id):
|
||
|
"""Begins a new branch in a cond section."""
|
||
|
assert section_id in self.cond_leaves
|
||
|
|
||
|
if section_id in self.cond_entry:
|
||
|
# Subsequent splits move back to the split point, and memorize the
|
||
|
# current leaves.
|
||
|
self.cond_leaves[section_id].append(self.leaves)
|
||
|
self.leaves = self.cond_entry[section_id]
|
||
|
else:
|
||
|
# If this is the first time we split a section, just remember the split
|
||
|
# point.
|
||
|
self.cond_entry[section_id] = self.leaves
|
||
|
|
||
|
def exit_cond_section(self, section_id):
|
||
|
"""Exits a conditional section."""
|
||
|
for split in self.cond_leaves[section_id]:
|
||
|
self.leaves |= split
|
||
|
del self.cond_entry[section_id]
|
||
|
del self.cond_leaves[section_id]
|
||
|
|
||
|
def enter_finally_section(self, section_id):
|
||
|
"""Enters a finally section."""
|
||
|
# TODO(mdan): This, not the caller, should track the active sections.
|
||
|
self.finally_section_subgraphs[section_id] = [None, None]
|
||
|
if self.leaves:
|
||
|
self.finally_section_has_direct_flow[section_id] = True
|
||
|
else:
|
||
|
self.finally_section_has_direct_flow[section_id] = False
|
||
|
self.pending_finally_sections.add(section_id)
|
||
|
|
||
|
def exit_finally_section(self, section_id):
|
||
|
"""Exits a finally section."""
|
||
|
assert section_id not in self.pending_finally_sections, 'Empty finally?'
|
||
|
self.finally_section_subgraphs[section_id][1] = self.leaves
|
||
|
# If the guard can only be reached by a jump, then it will not flow
|
||
|
# into the statement that follows it.
|
||
|
if not self.finally_section_has_direct_flow[section_id]:
|
||
|
self.leaves = set()
|
||
|
del self.finally_section_has_direct_flow[section_id]
|
||
|
|
||
|
def build(self):
|
||
|
"""Returns the CFG accumulated so far and resets the builder.
|
||
|
|
||
|
Returns:
|
||
|
Graph
|
||
|
"""
|
||
|
# Freeze the nodes.
|
||
|
for node in self.node_index.values():
|
||
|
node.freeze()
|
||
|
|
||
|
# Build the statement edges.
|
||
|
stmt_next = {}
|
||
|
stmt_prev = {}
|
||
|
for node, _ in self.forward_edges:
|
||
|
for stmt in self.owners[node]:
|
||
|
if stmt not in stmt_next:
|
||
|
stmt_next[stmt] = set()
|
||
|
if stmt not in stmt_prev:
|
||
|
stmt_prev[stmt] = set()
|
||
|
for first, second in self.forward_edges:
|
||
|
stmts_exited = self.owners[first] - self.owners[second]
|
||
|
for stmt in stmts_exited:
|
||
|
stmt_next[stmt].add(second)
|
||
|
stmts_entered = self.owners[second] - self.owners[first]
|
||
|
for stmt in stmts_entered:
|
||
|
stmt_prev[stmt].add(first)
|
||
|
for stmt in stmt_next:
|
||
|
stmt_next[stmt] = frozenset(stmt_next[stmt])
|
||
|
for stmt in stmt_prev:
|
||
|
stmt_prev[stmt] = frozenset(stmt_prev[stmt])
|
||
|
|
||
|
# Construct the final graph object.
|
||
|
result = Graph(
|
||
|
entry=self.head,
|
||
|
exit=self.leaves,
|
||
|
error=self.errors,
|
||
|
index=self.node_index,
|
||
|
stmt_prev=stmt_prev,
|
||
|
stmt_next=stmt_next)
|
||
|
|
||
|
# Reset the state.
|
||
|
self.reset()
|
||
|
|
||
|
return result
|
||
|
|
||
|
|
||
|
class AstToCfg(gast.NodeVisitor):
|
||
|
"""Converts an AST to CFGs.
|
||
|
|
||
|
A separate CFG will be constructed for each function.
|
||
|
"""
|
||
|
|
||
|
def __init__(self):
|
||
|
super(AstToCfg, self).__init__()
|
||
|
|
||
|
self.builder_stack = []
|
||
|
self.builder = None
|
||
|
self.cfgs = {}
|
||
|
|
||
|
self.lexical_scopes = []
|
||
|
|
||
|
def _enter_lexical_scope(self, node):
|
||
|
self.lexical_scopes.append(node)
|
||
|
|
||
|
def _exit_lexical_scope(self, node):
|
||
|
leaving_node = self.lexical_scopes.pop()
|
||
|
assert node == leaving_node
|
||
|
|
||
|
def _get_enclosing_scopes(self, include, stop_at):
|
||
|
included = []
|
||
|
for node in reversed(self.lexical_scopes):
|
||
|
if isinstance(node, include):
|
||
|
included.append(node)
|
||
|
if isinstance(node, stop_at):
|
||
|
return node, included
|
||
|
return None, included
|
||
|
|
||
|
def _process_basic_statement(self, node):
|
||
|
self.generic_visit(node)
|
||
|
self.builder.add_ordinary_node(node)
|
||
|
|
||
|
def _process_exit_statement(self, node, *exits_nodes_of_type):
|
||
|
# Note: this is safe because we process functions separately.
|
||
|
try_node, guards = self._get_enclosing_scopes(
|
||
|
include=(gast.Try,),
|
||
|
stop_at=tuple(exits_nodes_of_type),
|
||
|
)
|
||
|
if try_node is None:
|
||
|
raise ValueError(
|
||
|
'%s that is not enclosed by any of %s' % (node, exits_nodes_of_type))
|
||
|
self.builder.add_exit_node(node, try_node, guards)
|
||
|
|
||
|
def _process_continue_statement(self, node, *loops_to_nodes_of_type):
|
||
|
# Note: this is safe because we process functions separately.
|
||
|
try_node, guards = self._get_enclosing_scopes(
|
||
|
include=(gast.Try,),
|
||
|
stop_at=tuple(loops_to_nodes_of_type),
|
||
|
)
|
||
|
if try_node is None:
|
||
|
raise ValueError('%s that is not enclosed by any of %s' %
|
||
|
(node, loops_to_nodes_of_type))
|
||
|
self.builder.add_continue_node(node, try_node, guards)
|
||
|
|
||
|
def visit_FunctionDef(self, node):
|
||
|
# We also keep the FunctionDef node in the CFG. This allows us to determine
|
||
|
# things like reaching definitions via closure. Note that the function body
|
||
|
# will be stored in a separate graph, because function definitions are not
|
||
|
# the same as function calls.
|
||
|
if self.builder is not None:
|
||
|
self.builder.add_ordinary_node(node)
|
||
|
|
||
|
self.builder_stack.append(self.builder)
|
||
|
self.builder = GraphBuilder(node)
|
||
|
|
||
|
self._enter_lexical_scope(node)
|
||
|
self.builder.enter_section(node)
|
||
|
|
||
|
self._process_basic_statement(node.args)
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
self.builder.exit_section(node)
|
||
|
self._exit_lexical_scope(node)
|
||
|
|
||
|
self.cfgs[node] = self.builder.build()
|
||
|
self.builder = self.builder_stack.pop()
|
||
|
|
||
|
def visit_Lambda(self, node):
|
||
|
# TODO(mdan): Treat like FunctionDef? That would be a separate CFG.
|
||
|
raise NotImplementedError()
|
||
|
|
||
|
def visit_Return(self, node):
|
||
|
self._process_exit_statement(node, gast.FunctionDef)
|
||
|
|
||
|
def visit_Expr(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_Assign(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_AnnAssign(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_AugAssign(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_Print(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_Raise(self, node):
|
||
|
try_node, guards = self._get_enclosing_scopes(
|
||
|
include=(gast.Try,),
|
||
|
stop_at=(gast.FunctionDef,),
|
||
|
)
|
||
|
if try_node is None:
|
||
|
raise ValueError('%s that is not enclosed by any FunctionDef' % node)
|
||
|
self.builder.add_error_node(node, try_node, guards)
|
||
|
|
||
|
def visit_Assert(self, node):
|
||
|
# Ignoring the effect of exceptions.
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_Delete(self, node):
|
||
|
self._process_basic_statement(node)
|
||
|
|
||
|
def visit_If(self, node):
|
||
|
# No need to track ifs as lexical scopes, for now.
|
||
|
# Lexical scopes are generally tracked in order to be able to resolve the
|
||
|
# targets of jump statements like break/continue/etc. Since there is no
|
||
|
# statement that can interrupt a conditional, we don't need to track their
|
||
|
# lexical scope. That may change in the future.
|
||
|
self.builder.begin_statement(node)
|
||
|
|
||
|
self.builder.enter_cond_section(node)
|
||
|
self._process_basic_statement(node.test)
|
||
|
|
||
|
self.builder.new_cond_branch(node)
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
self.builder.new_cond_branch(node)
|
||
|
for stmt in node.orelse:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
self.builder.exit_cond_section(node)
|
||
|
self.builder.end_statement(node)
|
||
|
|
||
|
def visit_While(self, node):
|
||
|
self.builder.begin_statement(node)
|
||
|
self._enter_lexical_scope(node)
|
||
|
|
||
|
self.builder.enter_section(node)
|
||
|
|
||
|
self.builder.enter_loop_section(node, node.test)
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
self.builder.exit_loop_section(node)
|
||
|
|
||
|
# Note: although the orelse is technically part of the loop node,
|
||
|
# the statements inside it don't affect the loop itself. For example, a
|
||
|
# break in the loop's orelse will not affect the loop itself.
|
||
|
self._exit_lexical_scope(node)
|
||
|
|
||
|
for stmt in node.orelse:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
self.builder.exit_section(node)
|
||
|
self.builder.end_statement(node)
|
||
|
|
||
|
def visit_For(self, node):
|
||
|
self.builder.begin_statement(node)
|
||
|
self._enter_lexical_scope(node)
|
||
|
|
||
|
self.builder.enter_section(node)
|
||
|
|
||
|
# TODO(mdan): Strictly speaking, this should be node.target + node.iter.
|
||
|
# A blind dataflow analysis would have to process both node.target and
|
||
|
# node.iter to properly process read and write access.
|
||
|
self.builder.enter_loop_section(node, node.iter)
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
self.builder.exit_loop_section(node)
|
||
|
|
||
|
# Note: although the orelse is technically part of the loop node,
|
||
|
# they don't count as loop bodies. For example, a break in the loop's
|
||
|
# orelse will affect the parent loop, not the current one.
|
||
|
self._exit_lexical_scope(node)
|
||
|
|
||
|
for stmt in node.orelse:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
self.builder.exit_section(node)
|
||
|
self.builder.end_statement(node)
|
||
|
|
||
|
def visit_Break(self, node):
|
||
|
self._process_exit_statement(node, gast.While, gast.For)
|
||
|
|
||
|
def visit_Continue(self, node):
|
||
|
self._process_continue_statement(node, gast.While, gast.For)
|
||
|
|
||
|
def visit_Try(self, node):
|
||
|
self._enter_lexical_scope(node)
|
||
|
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
# Unlike loops, the orelse is a simple continuation of the body.
|
||
|
for stmt in node.orelse:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
if node.handlers:
|
||
|
# TODO(mdan): Should we still support bare try/except? Might be confusing.
|
||
|
raise NotImplementedError('exceptions are not yet supported')
|
||
|
|
||
|
self._exit_lexical_scope(node)
|
||
|
|
||
|
self.builder.enter_finally_section(node)
|
||
|
for stmt in node.finalbody:
|
||
|
self.visit(stmt)
|
||
|
self.builder.exit_finally_section(node)
|
||
|
|
||
|
def visit_With(self, node):
|
||
|
# TODO(mdan): Mark the context manager's exit call as exit guard.
|
||
|
for item in node.items:
|
||
|
self._process_basic_statement(item)
|
||
|
for stmt in node.body:
|
||
|
self.visit(stmt)
|
||
|
|
||
|
|
||
|
def build(node):
|
||
|
visitor = AstToCfg()
|
||
|
visitor.visit(node)
|
||
|
return visitor.cfgs
|