From f6d27f15fd5050ddce744cc6a1a560614af1a980 Mon Sep 17 00:00:00 2001
From: Kshitij Bansal <kshitijbansal@fb.com>
Date: Mon, 20 Dec 2021 17:11:19 -0800
Subject: [PATCH] Implement visit statement

Summary:
# This diff

Implements the visit/visit_statement for Match.

Reviewed By: grievejia

Differential Revision: D33243003

fbshipit-source-id: c57f17b23d834aaa2db96af69d20577e9e36e2c3
---
 source/ast/test/visitTest.ml | 42 +++++++++++++++++++++++++++++++++++-
 source/ast/visit.ml          | 34 +++++++++++++++++++++++++----
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/source/ast/test/visitTest.ml b/source/ast/test/visitTest.ml
index 50eaddc684e..3b91dc6dc05 100644
--- a/source/ast/test/visitTest.ml
+++ b/source/ast/test/visitTest.ml
@@ -138,6 +138,40 @@ let test_collect _ =
         +Statement.Expression (+Expression.Constant (Constant.Float 4.0));
         +Statement.Expression (+Expression.Constant (Constant.Float 3.0));
       ] );
+  assert_collect
+    [
+      +Statement.Match
+         {
+           Match.subject = +Expression.Constant (Constant.Integer 0);
+           cases =
+             [
+               {
+                 Match.Case.pattern = +Match.Pattern.MatchSingleton (Constant.Integer 2);
+                 guard = Some (+Expression.Constant (Constant.Integer 4));
+                 body = [];
+               };
+             ];
+         };
+    ]
+    ( [
+        +Expression.Constant (Constant.Integer 4);
+        +Expression.Constant (Constant.Integer 2);
+        +Expression.Constant (Constant.Integer 0);
+      ],
+      [
+        +Statement.Match
+           {
+             Match.subject = +Expression.Constant (Constant.Integer 0);
+             cases =
+               [
+                 {
+                   Match.Case.pattern = +Match.Pattern.MatchSingleton (Constant.Integer 2);
+                   guard = Some (+Expression.Constant (Constant.Integer 4));
+                   body = [];
+                 };
+               ];
+           };
+      ] );
 
   assert_collect
     [
@@ -385,6 +419,9 @@ let test_statement_visitor _ =
       | Import _ ->
           increment visited "import";
           visited
+      | Match _ ->
+          increment visited "match";
+          visited
       | Return _ ->
           increment visited "return";
           visited
@@ -409,9 +446,12 @@ let test_statement_visitor _ =
       if x < 3:
         import a
         c = 3
+      match x:
+        case _:
+          import b
   |}
   in
-  assert_counts source ["assign", 3; "return", 1; "import", 2];
+  assert_counts source ["assign", 3; "return", 1; "import", 3; "match", 1];
   ()
 
 
diff --git a/source/ast/visit.ml b/source/ast/visit.ml
index 372c16d52b1..002b691591c 100644
--- a/source/ast/visit.ml
+++ b/source/ast/visit.ml
@@ -185,6 +185,33 @@ module MakeNodeVisitor (Visitor : NodeVisitor) = struct
           visit_expression test;
           List.iter body ~f:visit_statement;
           List.iter orelse ~f:visit_statement
+      | Match { Match.subject; cases } ->
+          let rec visit_pattern { Node.value; location } =
+            match value with
+            | Match.Pattern.MatchAs { pattern; _ } -> Option.iter pattern ~f:visit_pattern
+            | MatchClass { patterns; keyword_patterns; _ } ->
+                List.iter patterns ~f:visit_pattern;
+                List.iter keyword_patterns ~f:visit_pattern
+            | MatchMapping { keys; patterns; _ } ->
+                List.iter keys ~f:visit_expression;
+                List.iter patterns ~f:visit_pattern
+            | MatchOr patterns
+            | MatchSequence patterns ->
+                List.iter patterns ~f:visit_pattern
+            | MatchSingleton constant ->
+                visit_expression { Node.value = Expression.Constant constant; location }
+            | MatchValue expression -> visit_expression expression
+            | MatchStar _
+            | MatchWildcard ->
+                ()
+          in
+          let visit_case { Match.Case.pattern; guard; body } =
+            visit_pattern pattern;
+            Option.iter guard ~f:visit_expression;
+            List.iter body ~f:visit_statement
+          in
+          visit_expression subject;
+          List.iter cases ~f:visit_case
       | Raise { Raise.expression; from } ->
           Option.iter ~f:visit_expression expression;
           Option.iter ~f:visit_expression from
@@ -212,8 +239,6 @@ module MakeNodeVisitor (Visitor : NodeVisitor) = struct
       | Import _
       | Nonlocal _
       | Global _
-      (* TODO(T107108860): Visitor for match statement. *)
-      | Match _
       | Pass
       | Continue
       | Break ->
@@ -280,8 +305,6 @@ module MakeStatementVisitor (Visitor : StatementVisitor) = struct
         | Expression _
         | Global _
         | Import _
-        (* TODO(T107108860): Visitor for match statement. *)
-        | Match _
         | Pass
         | Raise _
         | Return _
@@ -296,6 +319,9 @@ module MakeStatementVisitor (Visitor : StatementVisitor) = struct
         | While { While.body; orelse; _ } ->
             List.iter ~f:visit_statement body;
             List.iter ~f:visit_statement orelse
+        | Match { Match.cases; _ } ->
+            let visit_case { Match.Case.body; _ } = List.iter ~f:visit_statement body in
+            List.iter ~f:visit_case cases
         | Try { Try.body; handlers; orelse; finally } ->
             let visit_handler { Try.Handler.body; _ } = List.iter ~f:visit_statement body in
             List.iter ~f:visit_statement body;