xref: /aosp_15_r20/external/tensorflow/tensorflow/python/autograph/pyct/cfg.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Control flow graph (CFG) structure for Python AST representation.
16
17The CFG is a digraph with edges representing valid control flow. Each
18node is associated with exactly one AST node, but not all AST nodes may have
19a corresponding CFG counterpart.
20
21Once built, the CFG itself is immutable, but the values it holds need not be;
22they are usually annotated with information extracted by walking the graph.
23
24Tip: Use `Graph.as_dot` to visualize the CFG using any DOT viewer.
25
26Note: the CFG tries to include all code paths that MAY be taken, with a single
27notable exception:
28 * function calls do not generate edges corresponding to exceptions they may
29   raise (i.e. a function call in the middle of a block does not return or jump
30   to any except or finally block)
31TODO(mdan): Consider adding the edges above. They'd only add ~O(n) edges.
32TODO(mdan): Alternatively, consider adding an edge from try to all its excepts.
33"""
34
35# TODO(mdan): The notion of 'statements' below is inaccurate.
36# They should rather be called 'block statements', because they include
37# statements that may have a body, e.g. if and while.
38
39import collections
40import enum
41import weakref
42
43import astunparse
44import gast
45
46from tensorflow.python.autograph.pyct import anno
47
48
49class Node(object):
50  """A node in the CFG.
51
52  Although new instances of this class are mutable, the objects that a user
53  finds in the CFG are typically not.
54
55  The nodes represent edges in the CFG graph, and maintain pointers to allow
56  efficient walking in both forward and reverse order. The following property
57  holds for all nodes: "child in node.next" iff "node in child.prev".
58
59  Attributes:
60    next: FrozenSet[Node, ...], the nodes that follow this node, in control flow
61      order
62    prev: FrozenSet[Node, ...], the nodes that precede this node, in reverse
63      control flow order
64    ast_node: ast.AST, the AST node corresponding to this CFG node
65  """
66
67  def __init__(self, next_, prev, ast_node):
68    self.next = next_
69    self.prev = prev
70    self.ast_node = ast_node
71
72  def freeze(self):
73    self.next = frozenset(self.next)
74    # Assumption: All CFG nodes have identical life spans, because the graph
75    # owns them. Nodes should never be used outside the context of an existing
76    # graph.
77    self.prev = weakref.WeakSet(self.prev)
78
79  def __repr__(self):
80    if isinstance(self.ast_node, gast.FunctionDef):
81      return 'def %s' % self.ast_node.name
82    elif isinstance(self.ast_node, gast.ClassDef):
83      return 'class %s' % self.ast_node.name
84    elif isinstance(self.ast_node, gast.withitem):
85      # TODO(xjun): remove use of astunparse
86      return astunparse.unparse(self.ast_node.context_expr).strip()
87    return astunparse.unparse(self.ast_node).strip()
88
89
90class Graph(
91    collections.namedtuple(
92        'Graph',
93        ['entry', 'exit', 'error', 'index', 'stmt_prev', 'stmt_next'])):
94  """A Control Flow Graph.
95
96  The CFG maintains an index to allow looking up a CFG node by the AST node to
97  which it is associated. The index can also be enumerated in top-down, depth
98  first order.
99
100  Walking the graph in forward or reverse order is supported by double
101  parent-child links.
102
103  Note: the error nodes are not wired to their corresponding finally guards,
104  because these are shared, and wiring them would create a reverse path from
105  normal control flow into the error nodes, which we want to avoid.
106
107  The graph also maintains edges corresponding to higher level statements
108  like for-else loops. A node is considered successor of a statement if there
109  is an edge from a node that is lexically a child of that statement to a node
110  that is not. Statement predecessors are analogously defined.
111
112  Attributes:
113    entry: Node, the entry node
114    exit: FrozenSet[Node, ...], the exit nodes
115    error: FrozenSet[Node, ...], nodes that exit due to an explicitly raised
116      error (errors propagated from function calls are not accounted)
117    index: Dict[ast.Node, Node], mapping AST nodes to the respective CFG node
118    stmt_prev: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST nodes
119      to their predecessor CFG nodes
120    stmt_next: Dict[ast.Node, FrozenSet[Node, ...]], mapping statement AST nodes
121      to their successor CFG nodes
122  """
123
124  def __repr__(self):
125    return self.as_dot()
126
127  def as_dot(self):
128    """Print CFG in DOT format."""
129    result = 'digraph CFG {\n'
130    for node in self.index.values():
131      result += '  %s [label="%s"];\n' % (id(node), node)
132    for node in self.index.values():
133      for next_ in node.next:
134        result += '  %s -> %s;\n' % (id(node), id(next_))
135    result += '}'
136    return result
137
138
139class _WalkMode(enum.Enum):
140  FORWARD = 1
141  REVERSE = 2
142
143
144# TODO(mdan): Rename to DataFlowAnalyzer.
145# TODO(mdan): Consider specializations that use gen/kill/transfer abstractions.
146class GraphVisitor(object):
147  """Base class for a CFG visitors.
148
149  This implementation is not thread safe.
150
151  The visitor has some facilities to simplify dataflow analyses. In particular,
152  it allows revisiting the nodes at the decision of the subclass. This can be
153  used to visit the graph until the state reaches a fixed point.
154
155  For more details on dataflow analysis, see
156  https://www.seas.harvard.edu/courses/cs252/2011sp/slides/Lec02-Dataflow.pdf
157
158  Note: the literature generally suggests visiting successor nodes only when the
159  state of the current node changed, regardless of whether that successor has
160  ever been visited. This implementation visits every successor at least once.
161
162  Attributes:
163    graph: Graph
164    in_: Dict[Node, Any], stores node-keyed state during a visit
165    out: Dict[Node, Any], stores node-keyed state during a visit
166  """
167
168  def __init__(self, graph):
169    self.graph = graph
170    self.reset()
171
172  def init_state(self, node):
173    """State initialization function.
174
175    Optional to overload.
176
177    An in/out state slot will be created for each node in the graph. Subclasses
178    must overload this to control what that is initialized to.
179
180    Args:
181      node: Node
182    """
183    raise NotImplementedError('Subclasses must implement this.')
184
185  # TODO(mdan): Rename to flow?
186  def visit_node(self, node):
187    """Visitor function.
188
189    Args:
190      node: Node
191
192    Returns:
193      bool, whether the node should be revisited; subclasses can visit every
194          reachable node exactly once by always returning False
195    """
196    raise NotImplementedError('Subclasses must implement this.')
197
198  def reset(self):
199    self.in_ = {
200        node: self.init_state(node) for node in self.graph.index.values()
201    }
202    self.out = {
203        node: self.init_state(node) for node in self.graph.index.values()
204    }
205
206  def can_ignore(self, node):
207    """Returns True if the node can safely be assumed not to touch variables."""
208    ast_node = node.ast_node
209    if anno.hasanno(ast_node, anno.Basic.SKIP_PROCESSING):
210      return True
211    return isinstance(ast_node,
212                      (gast.Break, gast.Continue, gast.Raise, gast.Pass))
213
214  def _visit_internal(self, mode):
215    """Visits the CFG, breadth-first."""
216    assert mode in (_WalkMode.FORWARD, _WalkMode.REVERSE)
217    if mode == _WalkMode.FORWARD:
218      open_ = [self.graph.entry]
219    elif mode == _WalkMode.REVERSE:
220      open_ = list(self.graph.exit)
221    closed = set()
222
223    while open_:
224      node = open_.pop(0)
225      closed.add(node)
226
227      should_revisit = self.visit_node(node)
228
229      if mode == _WalkMode.FORWARD:
230        children = node.next
231      elif mode == _WalkMode.REVERSE:
232        children = node.prev
233
234      for next_ in children:
235        if should_revisit or next_ not in closed:
236          open_.append(next_)
237
238  def visit_forward(self):
239    self._visit_internal(_WalkMode.FORWARD)
240
241  def visit_reverse(self):
242    self._visit_internal(_WalkMode.REVERSE)
243
244
245class GraphBuilder(object):
246  """Builder that constructs a CFG from a given AST.
247
248  This GraphBuilder facilitates constructing the DAG that forms the CFG when
249  nodes
250  are supplied in lexical order (i.e., top-down, depth first). Under these
251  conditions, it supports building patterns found in typical structured
252  programs.
253
254  This builder ignores the flow generated by exceptions, which are assumed to
255  always be catastrophic and present purely for diagnostic purposes (e.g. to
256  print debug information). Statements like raise and try/catch sections are
257  allowed and will generate control flow edges, but ordinary statements are
258  assumed not to raise exceptions.
259
260  Finally sections are also correctly interleaved between break/continue/return
261  nodes and their subsequent statements.
262
263  Important concepts:
264   * nodes - nodes refer to CFG nodes; AST nodes are qualified explicitly
265   * leaf set - since the graph is constructed gradually, a leaf set maintains
266     the CFG nodes that will precede the node that the builder expects to
267     receive next; when an ordinary node is added, it is connected to the
268     existing leaves and it in turn becomes the new leaf
269   * jump nodes - nodes that should generate edges other than what
270     ordinary nodes would; these correspond to break, continue and return
271     statements
272   * sections - logical delimiters for subgraphs that require special
273     edges; there are various types of nodes, each admitting various
274     types of jump nodes; sections are identified by their corresponding AST
275     node
276  """
277
278  # TODO(mdan): Perhaps detail this in a markdown doc.
279  # TODO(mdan): Add exception support.
280
281  def __init__(self, parent_ast_node):
282    self.reset()
283    self.parent = parent_ast_node
284
285  def reset(self):
286    """Resets the state of this factory."""
287    self.head = None
288    self.errors = set()
289    self.node_index = {}
290
291    # TODO(mdan): Too many primitives. Use classes.
292    self.leaves = set()
293
294    # Note: This mechanism requires that nodes are added in lexical order (top
295    # to bottom, depth first).
296    self.active_stmts = set()
297    self.owners = {}  # type: Set[any]
298    self.forward_edges = set()  # type: Tuple[Node, Node] # (from, to)
299
300    self.finally_sections = {}
301    # Dict values represent (entry, exits)
302    self.finally_section_subgraphs = {
303    }  # type: Dict[ast.AST, Tuple[Node, Set[Node]]]
304    # Whether the guard section can be reached from the statement that precedes
305    # it.
306    self.finally_section_has_direct_flow = {}
307    # Finally sections that await their first node.
308    self.pending_finally_sections = set()
309
310    # Exit jumps keyed by the section they affect.
311    self.exits = {}
312
313    # The entry of loop sections, keyed by the section.
314    self.section_entry = {}
315    # Continue jumps keyed by the section they affect.
316    self.continues = {}
317
318    # Raise jumps keyed by the except section guarding them.
319    self.raises = {}
320
321    # The entry of conditional sections, keyed by the section.
322    self.cond_entry = {}
323    # Lists of leaf nodes corresponding to each branch in the section.
324    self.cond_leaves = {}
325
326  def _connect_nodes(self, first, second):
327    """Connects nodes to signify that control flows from first to second.
328
329    Args:
330      first: Union[Set[Node, ...], Node]
331      second: Node
332    """
333    if isinstance(first, Node):
334      first.next.add(second)
335      second.prev.add(first)
336      self.forward_edges.add((first, second))
337    else:
338      for node in first:
339        self._connect_nodes(node, second)
340
341  def _add_new_node(self, ast_node):
342    """Grows the graph by adding a CFG node following the current leaves."""
343    if ast_node in self.node_index:
344      raise ValueError('%s added twice' % ast_node)
345    # Assumption: All CFG nodes have identical life spans, because the graph
346    # owns them. Nodes should never be used outside the context of an existing
347    # graph.
348    node = Node(next_=set(), prev=weakref.WeakSet(), ast_node=ast_node)
349    self.node_index[ast_node] = node
350    self.owners[node] = frozenset(self.active_stmts)
351
352    if self.head is None:
353      self.head = node
354
355    for leaf in self.leaves:
356      self._connect_nodes(leaf, node)
357
358    # If any finally section awaits its first node, populate it.
359    for section_id in self.pending_finally_sections:
360      self.finally_section_subgraphs[section_id][0] = node
361    self.pending_finally_sections = set()
362
363    return node
364
365  def begin_statement(self, stmt):
366    """Marks the beginning of a statement.
367
368    Args:
369      stmt: Hashable, a key by which the statement can be identified in the
370        CFG's stmt_prev and stmt_next attributes
371    """
372    self.active_stmts.add(stmt)
373
374  def end_statement(self, stmt):
375    """Marks the end of a statement.
376
377    Args:
378      stmt: Hashable, a key by which the statement can be identified in the
379        CFG's stmt_prev and stmt_next attributes; must match a key previously
380        passed to begin_statement.
381    """
382    self.active_stmts.remove(stmt)
383
384  def add_ordinary_node(self, ast_node):
385    """Grows the graph by adding an ordinary CFG node.
386
387    Ordinary nodes are followed by the next node, in lexical order, that is,
388    they become the new leaf set.
389
390    Args:
391      ast_node: ast.AST
392
393    Returns:
394      Node
395    """
396    node = self._add_new_node(ast_node)
397    self.leaves = set((node,))
398    return node
399
400  def _add_jump_node(self, ast_node, guards):
401    """Grows the graph by adding a jump node.
402
403    Jump nodes are added to the current leaf set, and the leaf set becomes
404    empty. If the jump node is the last in a cond section, then it may be added
405    back to the leaf set by a separate mechanism.
406
407    Args:
408      ast_node: ast.AST
409      guards: Tuple[ast.AST, ...], the finally sections active for this node
410
411    Returns:
412      Node
413    """
414    node = self._add_new_node(ast_node)
415    self.leaves = set()
416    # The guards themselves may not yet be complete, and will be wired later.
417    self.finally_sections[node] = guards
418    return node
419
420  def _connect_jump_to_finally_sections(self, node):
421    """Connects a jump node to the finally sections protecting it."""
422    cursor = set((node,))
423    if node not in self.finally_sections:
424      return cursor
425    for guard_section_id in self.finally_sections[node]:
426      guard_begin, guard_ends = self.finally_section_subgraphs[guard_section_id]
427      self._connect_nodes(cursor, guard_begin)
428      cursor = guard_ends
429    del self.finally_sections[node]
430    # TODO(mdan): Should garbage-collect finally_section_subgraphs.
431    return cursor
432
433  def add_exit_node(self, ast_node, section_id, guards):
434    """Grows the graph by adding an exit node.
435
436    This node becomes an exit for the current section.
437
438    Args:
439      ast_node: ast.AST
440      section_id: Hashable, the node for which ast_node should be considered to
441        be an exit node
442      guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
443
444    Returns:
445      Node
446    """
447    node = self._add_jump_node(ast_node, guards)
448    self.exits[section_id].add(node)
449    return node
450
451  def add_continue_node(self, ast_node, section_id, guards):
452    """Grows the graph by adding a reentry node.
453
454    This node causes control flow to go back to the loop section's entry.
455
456    Args:
457      ast_node: ast.AST
458      section_id: Hashable, the node for which ast_node should be considered to
459        be an exit node
460      guards: Tuple[ast.AST, ...], the finally sections that guard ast_node
461    """
462    node = self._add_jump_node(ast_node, guards)
463    self.continues[section_id].add(node)
464
465  def connect_raise_node(self, node, except_guards):
466    """Adds extra connection between a raise node and containing except guards.
467
468    The node is a graph node, not an ast node.
469
470    Args:
471      node: Node
472      except_guards: Tuple[ast.AST, ...], the except sections that guard node
473    """
474    for guard in except_guards:
475      if guard in self.raises:
476        self.raises[guard].append(node)
477      else:
478        self.raises[guard] = [node]
479
480  def enter_section(self, section_id):
481    """Enters a regular section.
482
483    Regular sections admit exit jumps, which end the section.
484
485    Args:
486      section_id: Hashable, the same node that will be used in calls to the
487        ast_node arg passed to add_exit_node
488    """
489    assert section_id not in self.exits
490    self.exits[section_id] = set()
491
492  def exit_section(self, section_id):
493    """Exits a regular section."""
494
495    # Exits are jump nodes, which may be protected.
496    for exit_ in self.exits[section_id]:
497      self.leaves |= self._connect_jump_to_finally_sections(exit_)
498
499    del self.exits[section_id]
500
501  def enter_loop_section(self, section_id, entry_node):
502    """Enters a loop section.
503
504    Loop sections define an entry node. The end of the section always flows back
505    to the entry node. These admit continue jump nodes which also flow to the
506    entry node.
507
508    Args:
509      section_id: Hashable, the same node that will be used in calls to the
510        ast_node arg passed to add_continue_node
511      entry_node: ast.AST, the entry node into the loop (e.g. the test node for
512        while loops)
513    """
514    assert section_id not in self.section_entry
515    assert section_id not in self.continues
516    self.continues[section_id] = set()
517    node = self.add_ordinary_node(entry_node)
518    self.section_entry[section_id] = node
519
520  def exit_loop_section(self, section_id):
521    """Exits a loop section."""
522    self._connect_nodes(self.leaves, self.section_entry[section_id])
523
524    # continues are jump nodes, which may be protected.
525    for reentry in self.continues[section_id]:
526      guard_ends = self._connect_jump_to_finally_sections(reentry)
527      self._connect_nodes(guard_ends, self.section_entry[section_id])
528
529    # Loop nodes always loop back.
530    self.leaves = set((self.section_entry[section_id],))
531
532    del self.continues[section_id]
533    del self.section_entry[section_id]
534
535  def enter_cond_section(self, section_id):
536    """Enters a conditional section.
537
538    Conditional sections define an entry node, and one or more branches.
539
540    Args:
541      section_id: Hashable, the same node that will be used in calls to the
542        section_id arg passed to new_cond_branch
543    """
544
545    assert section_id not in self.cond_entry
546    assert section_id not in self.cond_leaves
547    self.cond_leaves[section_id] = []
548
549  def new_cond_branch(self, section_id):
550    """Begins a new branch in a cond section."""
551    assert section_id in self.cond_leaves
552
553    if section_id in self.cond_entry:
554      # Subsequent splits move back to the split point, and memorize the
555      # current leaves.
556      self.cond_leaves[section_id].append(self.leaves)
557      self.leaves = self.cond_entry[section_id]
558    else:
559      # If this is the first time we split a section, just remember the split
560      # point.
561      self.cond_entry[section_id] = self.leaves
562
563  def exit_cond_section(self, section_id):
564    """Exits a conditional section."""
565    for split in self.cond_leaves[section_id]:
566      self.leaves |= split
567    del self.cond_entry[section_id]
568    del self.cond_leaves[section_id]
569
570  def enter_except_section(self, section_id):
571    """Enters an except section."""
572    if section_id in self.raises:
573      self.leaves.update(self.raises[section_id])
574
575  def enter_finally_section(self, section_id):
576    """Enters a finally section."""
577    # TODO(mdan): This, not the caller, should track the active sections.
578    self.finally_section_subgraphs[section_id] = [None, None]
579    if self.leaves:
580      self.finally_section_has_direct_flow[section_id] = True
581    else:
582      self.finally_section_has_direct_flow[section_id] = False
583    self.pending_finally_sections.add(section_id)
584
585  def exit_finally_section(self, section_id):
586    """Exits a finally section."""
587    assert section_id not in self.pending_finally_sections, 'Empty finally?'
588    self.finally_section_subgraphs[section_id][1] = self.leaves
589    # If the guard can only be reached by a jump, then it will not flow
590    # into the statement that follows it.
591    if not self.finally_section_has_direct_flow[section_id]:
592      self.leaves = set()
593    del self.finally_section_has_direct_flow[section_id]
594
595  def build(self):
596    """Returns the CFG accumulated so far and resets the builder.
597
598    Returns:
599      Graph
600    """
601    # Freeze the nodes.
602    for node in self.node_index.values():
603      node.freeze()
604
605    # Build the statement edges.
606    stmt_next = {}
607    stmt_prev = {}
608
609    for node in self.node_index.values():
610      for stmt in self.owners[node]:
611        if stmt not in stmt_prev:
612          stmt_prev[stmt] = set()
613        if stmt not in stmt_next:
614          stmt_next[stmt] = set()
615
616    for first, second in self.forward_edges:
617      stmts_exited = self.owners[first] - self.owners[second]
618      for stmt in stmts_exited:
619        stmt_next[stmt].add(second)
620      stmts_entered = self.owners[second] - self.owners[first]
621      for stmt in stmts_entered:
622        stmt_prev[stmt].add(first)
623    for stmt in stmt_next:
624      stmt_next[stmt] = frozenset(stmt_next[stmt])
625    for stmt in stmt_prev:
626      stmt_prev[stmt] = frozenset(stmt_prev[stmt])
627
628    # Construct the final graph object.
629    result = Graph(
630        entry=self.head,
631        exit=self.leaves,
632        error=self.errors,
633        index=self.node_index,
634        stmt_prev=stmt_prev,
635        stmt_next=stmt_next)
636
637    # Reset the state.
638    self.reset()
639
640    return result
641
642
643class AstToCfg(gast.NodeVisitor):
644  """Converts an AST to CFGs.
645
646  A separate CFG will be constructed for each function.
647  """
648
649  def __init__(self):
650    super(AstToCfg, self).__init__()
651
652    self.builder_stack = []
653    self.builder = None
654    self.cfgs = {}
655
656    self.lexical_scopes = []
657
658  def _enter_lexical_scope(self, node):
659    self.lexical_scopes.append(node)
660
661  def _exit_lexical_scope(self, node):
662    leaving_node = self.lexical_scopes.pop()
663    assert node == leaving_node
664
665  def _get_enclosing_finally_scopes(self, stop_at):
666    included = []
667    for node in reversed(self.lexical_scopes):
668      if isinstance(node, gast.Try) and node.finalbody:
669        included.append(node)
670      if isinstance(node, stop_at):
671        return node, included
672    return None, included
673
674  def _get_enclosing_except_scopes(self, stop_at):
675    included = []
676    for node in reversed(self.lexical_scopes):
677      if isinstance(node, gast.Try) and node.handlers:
678        included.extend(node.handlers)
679      if isinstance(node, stop_at):
680        break
681    return included
682
683  def _process_basic_statement(self, node):
684    self.generic_visit(node)
685    self.builder.add_ordinary_node(node)
686
687  def _process_exit_statement(self,
688                              node,
689                              exits_nodes_of_type,
690                              may_exit_via_except=False):
691    self.generic_visit(node)
692    # Note: this is safe because we process functions separately.
693    try_node, guards = self._get_enclosing_finally_scopes(exits_nodes_of_type)
694    assert try_node is not None, '{} that is not enclosed by any of {}'.format(
695        node, exits_nodes_of_type)
696
697    node = self.builder.add_exit_node(node, try_node, guards)
698
699    if may_exit_via_except:
700      except_guards = self._get_enclosing_except_scopes(exits_nodes_of_type)
701      self.builder.connect_raise_node(node, except_guards)
702
703  def _process_continue_statement(self, node, *loops_to_nodes_of_type):
704    # Note: this is safe because we process functions separately.
705    try_node, guards = self._get_enclosing_finally_scopes(
706        tuple(loops_to_nodes_of_type))
707    if try_node is None:
708      raise ValueError('%s that is not enclosed by any of %s' %
709                       (node, loops_to_nodes_of_type))
710    self.builder.add_continue_node(node, try_node, guards)
711
712  def visit_ClassDef(self, node):
713    # We also keep the ClassDef node in the CFG, since it technically is a
714    # statement.
715    # For example, this is legal and allows executing user code:
716    #
717    #   class Foo(bar()):
718    #     pass
719    #
720    # It also has a scope:
721    #
722    #   class Bar(object):
723    #     a = 1
724    if self.builder is None:
725      self.generic_visit(node)
726      return
727
728    self.builder.add_ordinary_node(node)
729
730    self.builder_stack.append(self.builder)
731    self.builder = GraphBuilder(node)
732    self._enter_lexical_scope(node)
733
734    self._process_basic_statement(node)
735
736    self._exit_lexical_scope(node)
737    # TODO(mdan): Track the CFG local to the class definition as well?
738    self.builder = self.builder_stack.pop()
739
740  def _process_function_def(self, node, is_lambda):
741    # The function body is stored in a separate graph, because function
742    # definitions have effects very different from function calls.
743    if self.builder is not None:
744      self.builder.add_ordinary_node(node)
745
746    self.builder_stack.append(self.builder)
747    self.builder = GraphBuilder(node)
748
749    self._enter_lexical_scope(node)
750    self.builder.enter_section(node)
751
752    self._process_basic_statement(node.args)
753    if is_lambda:
754      self._process_exit_statement(node.body, (gast.Lambda,))
755    else:
756      for stmt in node.body:
757        self.visit(stmt)
758
759    self.builder.exit_section(node)
760    self._exit_lexical_scope(node)
761
762    self.cfgs[node] = self.builder.build()
763    self.builder = self.builder_stack.pop()
764
765  def visit_FunctionDef(self, node):
766    self._process_function_def(node, is_lambda=False)
767
768  def visit_Lambda(self, node):
769    self._process_function_def(node, is_lambda=True)
770
771  def visit_Return(self, node):
772    self._process_exit_statement(node, (gast.FunctionDef,))
773
774  def visit_Import(self, node):
775    self._process_basic_statement(node)
776
777  def visit_ImportFrom(self, node):
778    self._process_basic_statement(node)
779
780  def visit_Expr(self, node):
781    self._process_basic_statement(node)
782
783  def visit_Assign(self, node):
784    self._process_basic_statement(node)
785
786  def visit_AnnAssign(self, node):
787    self._process_basic_statement(node)
788
789  def visit_AugAssign(self, node):
790    self._process_basic_statement(node)
791
792  def visit_Pass(self, node):
793    self._process_basic_statement(node)
794
795  def visit_Global(self, node):
796    self._process_basic_statement(node)
797
798  def visit_Nonlocal(self, node):
799    self._process_basic_statement(node)
800
801  def visit_Print(self, node):
802    self._process_basic_statement(node)
803
804  def visit_Raise(self, node):
805    self._process_exit_statement(
806        node, (gast.FunctionDef,), may_exit_via_except=True)
807    self.builder.errors.add(node)
808
809  def visit_Assert(self, node):
810    # Ignoring the effect of exceptions.
811    self._process_basic_statement(node)
812
813  def visit_Delete(self, node):
814    self._process_basic_statement(node)
815
816  def visit_If(self, node):
817    # No need to track ifs as lexical scopes, for now.
818    # Lexical scopes are generally tracked in order to be able to resolve the
819    # targets of jump statements like break/continue/etc. Since there is no
820    # statement that can interrupt a conditional, we don't need to track their
821    # lexical scope. That may change in the future.
822    self.builder.begin_statement(node)
823
824    self.builder.enter_cond_section(node)
825    self._process_basic_statement(node.test)
826
827    self.builder.new_cond_branch(node)
828    for stmt in node.body:
829      self.visit(stmt)
830
831    self.builder.new_cond_branch(node)
832    for stmt in node.orelse:
833      self.visit(stmt)
834
835    self.builder.exit_cond_section(node)
836    self.builder.end_statement(node)
837
838  def visit_While(self, node):
839    self.builder.begin_statement(node)
840    self._enter_lexical_scope(node)
841
842    self.builder.enter_section(node)
843
844    self.generic_visit(node.test)
845    self.builder.enter_loop_section(node, node.test)
846    for stmt in node.body:
847      self.visit(stmt)
848    self.builder.exit_loop_section(node)
849
850    # Note: although the orelse is technically part of the loop node,
851    # the statements inside it don't affect the loop itself. For example, a
852    # break in the loop's orelse will not affect the loop itself.
853    self._exit_lexical_scope(node)
854
855    for stmt in node.orelse:
856      self.visit(stmt)
857
858    self.builder.exit_section(node)
859    self.builder.end_statement(node)
860
861  def visit_For(self, node):
862    self.builder.begin_statement(node)
863    self._enter_lexical_scope(node)
864
865    self.builder.enter_section(node)
866
867    # Note: Strictly speaking, this should be node.target + node.iter.
868    # However, the activity analysis accounts for this inconsistency,
869    # so dataflow analysis produces the correct values.
870    self.generic_visit(node.iter)
871    self.builder.enter_loop_section(node, node.iter)
872    # Also include the "extra loop test" annotation, to capture things like the
873    # control variable for return and break in for loops.
874    if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
875      self._process_basic_statement(
876          anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST))
877    for stmt in node.body:
878      self.visit(stmt)
879    self.builder.exit_loop_section(node)
880
881    # Note: although the orelse is technically part of the loop node,
882    # they don't count as loop bodies.  For example, a break in the loop's
883    # orelse will affect the parent loop, not the current one.
884    self._exit_lexical_scope(node)
885
886    for stmt in node.orelse:
887      self.visit(stmt)
888
889    self.builder.exit_section(node)
890    self.builder.end_statement(node)
891
892  def visit_Break(self, node):
893    self._process_exit_statement(node, (
894        gast.While,
895        gast.For,
896    ))
897
898  def visit_Continue(self, node):
899    self._process_continue_statement(node, (
900        gast.While,
901        gast.For,
902    ))
903
904  def visit_ExceptHandler(self, node):
905    self.builder.begin_statement(node)
906    self.builder.enter_except_section(node)
907
908    if node.type is not None:
909      self.visit(node.type)
910    if node.name is not None:
911      self.visit(node.name)
912
913    for stmt in node.body:
914      self.visit(stmt)
915
916    self.builder.end_statement(node)
917
918  def visit_Try(self, node):
919    self.builder.begin_statement(node)
920    self._enter_lexical_scope(node)
921
922    # Note: the current simplification is that the try block fully executes
923    # regardless of whether an exception triggers or not. This is consistent
924    # with blocks free of try/except, which also don't account for the
925    # possibility of an exception being raised mid-block.
926
927    for stmt in node.body:
928      self.visit(stmt)
929    # The orelse is an optional continuation of the body.
930    if node.orelse:
931      block_representative = node.orelse[0]
932      self.builder.enter_cond_section(block_representative)
933      self.builder.new_cond_branch(block_representative)
934      for stmt in node.orelse:
935        self.visit(stmt)
936      self.builder.new_cond_branch(block_representative)
937      self.builder.exit_cond_section(block_representative)
938
939    self._exit_lexical_scope(node)
940
941    if node.handlers:
942      # Using node would be inconsistent. Using the first handler node is also
943      # inconsistent, but less so.
944      block_representative = node.handlers[0]
945      self.builder.enter_cond_section(block_representative)
946      for block in node.handlers:
947        self.builder.new_cond_branch(block_representative)
948        self.visit(block)
949      self.builder.new_cond_branch(block_representative)
950      self.builder.exit_cond_section(block_representative)
951
952    if node.finalbody:
953      self.builder.enter_finally_section(node)
954      for stmt in node.finalbody:
955        self.visit(stmt)
956      self.builder.exit_finally_section(node)
957
958    self.builder.end_statement(node)
959
960  def visit_With(self, node):
961    # TODO(mdan): Mark the context manager's exit call as exit guard.
962    for item in node.items:
963      self._process_basic_statement(item)
964    for stmt in node.body:
965      self.visit(stmt)
966
967
968def build(node):
969  visitor = AstToCfg()
970  visitor.visit(node)
971  return visitor.cfgs
972