From e3683d8fbcd07aefb8fe941a5080a7c68e375f39 Mon Sep 17 00:00:00 2001
From: Rebecca Chen <rechen@google.com>
Date: Thu, 20 Oct 2022 15:21:28 -0700
Subject: [PATCH] Copy parser._Matches to parser_libcst._Matches.

---
 pytype/directors/directors.py     |  1 -
 pytype/directors/parser_libcst.py | 23 +++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/pytype/directors/directors.py b/pytype/directors/directors.py
index a6fd1fc4e..5d1b62f6d 100644
--- a/pytype/directors/directors.py
+++ b/pytype/directors/directors.py
@@ -287,7 +287,6 @@ def __init__(self, src_tree, errorlog, filename, disable, code):
     self._function_ranges = _BlockRanges({})
     # Parse the source code for directives.
     self._parse_src_tree(src_tree, code)
-    self._matches = None
 
   @property
   def type_comments(self):
diff --git a/pytype/directors/parser_libcst.py b/pytype/directors/parser_libcst.py
index 17026af8b..15b4204eb 100644
--- a/pytype/directors/parser_libcst.py
+++ b/pytype/directors/parser_libcst.py
@@ -81,6 +81,28 @@ def __repr__(self):
     """
 
 
+class _Matches:
+  """Tracks branches of match statements."""
+
+  def __init__(self):
+    self.start_to_end = {}
+    self.end_to_starts = collections.defaultdict(list)
+    self.match_cases = {}
+
+  def add_match(self, start, end, cases):
+    self.start_to_end[start] = end
+    self.end_to_starts[end].append(start)
+    for case_start, case_end in cases:
+      for i in range(case_start, case_end + 1):
+        self.match_cases[i] = start
+
+  def __repr__(self):
+    return f"""
+      Matches: {sorted(self.start_to_end.items())}
+      Cases: {self.match_cases}
+    """
+
+
 class _ParseVisitor(libcst.CSTVisitor):
   """Visitor for parsing a source tree.
 
@@ -110,6 +132,7 @@ def __init__(self):
     self.defs_start = None
     self.function_ranges = {}
     self.block_returns = _BlockReturns()
+    self.matches = _Matches()
 
   def _get_containing_groups(self, start_line, end_line=None):
     """Get _StructuredComment groups that fully contain the given line range."""