From 28a3c35ae36c996a9a64b516837b56d369ea0259 Mon Sep 17 00:00:00 2001
From: Danny Yang <yangdanny@meta.com>
Date: Tue, 14 Jan 2025 15:39:56 -0800
Subject: [PATCH] improve refinement for match

Summary:
When we generate conditions on an index of a sequence, if the sequence is a tuple literal and the index is in-bounds, we can generate the condition on the appropriate sub-expression within the tuple literal, instead of synthesizing a new subscript expression.

One gap that I discovered but didn't fix with this diff is for MatchAs, for example `match subject: case (1, 2) as z:`, which does not give the correct type to z. This is because for MatchAs we generate two refinement conditions, `z == subject && z == (1, 2)`, where `subject: tuple[object, object]`. To combine the refinements we `meet` the two types, and the result is `tuple[object, object]`, even though `tuple[1, 2]` is narrower. I think the desired behavior should be intersection, but probably not worth doing in pyre1.

Reviewed By: stroxler

Differential Revision: D68169119

fbshipit-source-id: 4d2867ec91414dcc8b1801739fca3d5179c1ea56
---
 source/analysis/cfg.ml                        | 13 +++++++++----
 source/analysis/test/integration/matchTest.ml | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/source/analysis/cfg.ml b/source/analysis/cfg.ml
index ea1c210cd2b..72529bc2ab3 100644
--- a/source/analysis/cfg.ml
+++ b/source/analysis/cfg.ml
@@ -55,10 +55,15 @@ module MatchTranslate = struct
 
 
   let create_subscript_index ~location ~sequence ~index =
-    create_subscript
-      ~location
-      ~container:sequence
-      ~key:(create_constant ~location (Constant.Integer index))
+    match sequence with
+    | { Node.value = Expression.Tuple elements; _ } when index >= 0 && index < List.length elements
+      ->
+        List.nth_exn elements index
+    | _ ->
+        create_subscript
+          ~location
+          ~container:sequence
+          ~key:(create_constant ~location (Constant.Integer index))
 
 
   let create_call ~location ~callee ~arguments =
diff --git a/source/analysis/test/integration/matchTest.ml b/source/analysis/test/integration/matchTest.ml
index ff21169ea11..04d73beec90 100644
--- a/source/analysis/test/integration/matchTest.ml
+++ b/source/analysis/test/integration/matchTest.ml
@@ -12,6 +12,24 @@ let test_simple =
   test_list
     [
       labeled_test_case __FUNCTION__ __LINE__
+      @@ assert_type_errors
+           {|
+      def foo(x: object, y: object) -> None:
+        match (x, y):
+          case (1, 2):
+            reveal_type(x)
+            reveal_type(y)
+          case (1 as x2, 2 as y2):
+            reveal_type(x2)
+            reveal_type(y2)
+    |}
+           [
+             "Revealed type [-1]: Revealed type for `x` is `typing_extensions.Literal[1]`.";
+             "Revealed type [-1]: Revealed type for `y` is `typing_extensions.Literal[2]`.";
+             "Revealed type [-1]: Revealed type for `x2` is `typing_extensions.Literal[1]`.";
+             "Revealed type [-1]: Revealed type for `y2` is `typing_extensions.Literal[2]`.";
+           ];
+      labeled_test_case __FUNCTION__ __LINE__
       @@ assert_type_errors
            {|
       def foo(x: int) -> None: