[cfe] Support pattern entry-to-element conversion in type inference

Part of #49749

Change-Id: I5a227763225eda6a8154591fa35aa7cc66c0050a
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/285320
Reviewed-by: Johnni Winther <[email protected]>
Commit-Queue: Chloe Stefantsova <[email protected]>

Author: chloestefantsova

pkg/front_end/lib/src/fasta/kernel/body_builder.dart

@@ -6692,16 +6692,22 @@ class BodyBuilder extends StackListenerImpl
         node = forest.createIfMapEntry(
             offsetForToken(ifToken), condition.expression, entry);
       } else {
-        node = forest.createIfCaseMapEntry(
-            offsetForToken(ifToken), condition.expression, patternGuard, entry);
+        node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
+            prelude: [],
+            expression: condition.expression,
+            patternGuard: patternGuard,
+            then: entry);
       }
     } else {
       if (patternGuard == null) {
         node = forest.createIfElement(
             offsetForToken(ifToken), condition.expression, toValue(entry));
       } else {
         node = forest.createIfCaseElement(offsetForToken(ifToken),
-            condition.expression, patternGuard, toValue(entry));
+            prelude: [],
+            expression: condition.expression,
+            patternGuard: patternGuard,
+            then: toValue(entry));
       }
     }
     push(node);
@@ -6747,7 +6753,11 @@ class BodyBuilder extends StackListenerImpl
               condition.expression, thenEntry, elseEntry);
         } else {
           node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
-              condition.expression, patternGuard, thenEntry, elseEntry);
+              prelude: [],
+              expression: condition.expression,
+              patternGuard: patternGuard,
+              then: thenEntry,
+              otherwise: elseEntry);
         }
       } else if (elseEntry is ControlFlowElement) {
         MapLiteralEntry? elseMapEntry = elseEntry
@@ -6758,7 +6768,11 @@ class BodyBuilder extends StackListenerImpl
                 condition.expression, thenEntry, elseMapEntry);
           } else {
             node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
-                condition.expression, patternGuard, thenEntry, elseMapEntry);
+                prelude: [],
+                expression: condition.expression,
+                patternGuard: patternGuard,
+                then: thenEntry,
+                otherwise: elseMapEntry);
           }
         } else {
           int offset = elseEntry.fileOffset;
@@ -6788,7 +6802,11 @@ class BodyBuilder extends StackListenerImpl
                 condition.expression, thenMapEntry, elseEntry);
           } else {
             node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
-                condition.expression, patternGuard, thenMapEntry, elseEntry);
+                prelude: [],
+                expression: condition.expression,
+                patternGuard: patternGuard,
+                then: thenMapEntry,
+                otherwise: elseEntry);
           }
         } else {
           int offset = thenEntry.fileOffset;
@@ -6813,12 +6831,12 @@ class BodyBuilder extends StackListenerImpl
         node = forest.createIfElement(offsetForToken(ifToken),
             condition.expression, toValue(thenEntry), toValue(elseEntry));
       } else {
-        node = forest.createIfCaseElement(
-            offsetForToken(ifToken),
-            condition.expression,
-            condition.patternGuard!,
-            toValue(thenEntry),
-            toValue(elseEntry));
+        node = forest.createIfCaseElement(offsetForToken(ifToken),
+            prelude: [],
+            expression: condition.expression,
+            patternGuard: condition.patternGuard!,
+            then: toValue(thenEntry),
+            otherwise: toValue(elseEntry));
       }
     }
     push(node);

pkg/front_end/lib/src/fasta/kernel/collections.dart

@@ -837,13 +837,42 @@ Expression convertToElement(
     onConvertMapEntry(entry, result);
     return result;
   }
-  if (entry is ForMapEntry) {
-    ForElement result = new ForElement(entry.variables, entry.condition,
-        entry.updates, convertToElement(entry.body, helper, onConvertMapEntry))
+  if (entry is IfCaseMapEntry) {
+    IfCaseElement result = new IfCaseElement(
+        prelude: entry.prelude,
+        expression: entry.expression,
+        patternGuard: entry.patternGuard,
+        then: convertToElement(entry.then, helper, onConvertMapEntry),
+        otherwise: entry.otherwise == null
+            ? null
+            : convertToElement(entry.otherwise!, helper, onConvertMapEntry))
       ..fileOffset = entry.fileOffset;
     onConvertMapEntry(entry, result);
     return result;
   }
+  if (entry is ForMapEntry) {
+    if (entry is PatternForMapEntry) {
+      PatternForElement result = new PatternForElement(
+          patternVariableDeclaration: entry.patternVariableDeclaration,
+          prelude: entry.prelude,
+          variables: entry.variables,
+          condition: entry.condition,
+          updates: entry.updates,
+          body: convertToElement(entry.body, helper, onConvertMapEntry))
+        ..fileOffset = entry.fileOffset;
+      onConvertMapEntry(entry, result);
+      return result;
+    } else {
+      ForElement result = new ForElement(
+          entry.variables,
+          entry.condition,
+          entry.updates,
+          convertToElement(entry.body, helper, onConvertMapEntry))
+        ..fileOffset = entry.fileOffset;
+      onConvertMapEntry(entry, result);
+      return result;
+    }
+  }
   if (entry is ForInMapEntry) {
     ForInElement result = new ForInElement(
         entry.variable,
@@ -925,10 +954,11 @@ MapLiteralEntry convertToMapEntry(Expression element, InferenceHelper helper,
   }
   if (element is IfCaseElement) {
     IfCaseMapEntry result = new IfCaseMapEntry(
-        element.expression,
-        element.patternGuard,
-        convertToMapEntry(element.then, helper, onConvertElement),
-        element.otherwise == null
+        prelude: [],
+        expression: element.expression,
+        patternGuard: element.patternGuard,
+        then: convertToMapEntry(element.then, helper, onConvertElement),
+        otherwise: element.otherwise == null
             ? null
             : convertToMapEntry(element.otherwise!, helper, onConvertElement))
       ..fileOffset = element.fileOffset;

pkg/front_end/lib/src/fasta/kernel/forest.dart

@@ -272,12 +272,20 @@ class Forest {
     return new IfElement(condition, then, otherwise)..fileOffset = fileOffset;
   }
 
-  Expression createIfCaseElement(int fileOffset, Expression expression,
-      PatternGuard patternGuard, Expression then,
-      [Expression? otherwise]) {
+  Expression createIfCaseElement(int fileOffset,
+      {required List<Statement> prelude,
+      required Expression expression,
+      required PatternGuard patternGuard,
+      required Expression then,
+      Expression? otherwise}) {
     // ignore: unnecessary_null_comparison
     assert(fileOffset != null);
-    return new IfCaseElement(expression, patternGuard, then, otherwise)
+    return new IfCaseElement(
+        prelude: prelude,
+        expression: expression,
+        patternGuard: patternGuard,
+        then: then,
+        otherwise: otherwise)
       ..fileOffset = fileOffset;
   }
 
@@ -289,12 +297,20 @@ class Forest {
     return new IfMapEntry(condition, then, otherwise)..fileOffset = fileOffset;
   }
 
-  MapLiteralEntry createIfCaseMapEntry(int fileOffset, Expression expression,
-      PatternGuard patternGuard, MapLiteralEntry then,
-      [MapLiteralEntry? otherwise]) {
+  MapLiteralEntry createIfCaseMapEntry(int fileOffset,
+      {required List<Statement> prelude,
+      required Expression expression,
+      required PatternGuard patternGuard,
+      required MapLiteralEntry then,
+      MapLiteralEntry? otherwise}) {
     // ignore: unnecessary_null_comparison
     assert(fileOffset != null);
-    return new IfCaseMapEntry(expression, patternGuard, then, otherwise)
+    return new IfCaseMapEntry(
+        prelude: prelude,
+        expression: expression,
+        patternGuard: patternGuard,
+        then: then,
+        otherwise: otherwise)
       ..fileOffset = fileOffset;
   }
 

pkg/front_end/lib/src/fasta/kernel/internal_ast.dart

@@ -4813,9 +4813,15 @@ class IfCaseElement extends InternalExpression with ControlFlowElement {
   PatternGuard patternGuard;
   Expression then;
   Expression? otherwise;
-  late List<Statement> replacement;
+  List<Statement> prelude;
 
-  IfCaseElement(this.expression, this.patternGuard, this.then, this.otherwise) {
+  IfCaseElement(
+      {required this.prelude,
+      required this.expression,
+      required this.patternGuard,
+      required this.then,
+      this.otherwise}) {
+    setParents(prelude, this);
     expression.parent = this;
     patternGuard.parent = this;
     then.parent = this;
@@ -4861,9 +4867,13 @@ class IfCaseElement extends InternalExpression with ControlFlowElement {
       }
       if (otherwiseEntry == null) return null;
     }
-    IfCaseMapEntry result =
-        new IfCaseMapEntry(expression, patternGuard, thenEntry, otherwiseEntry)
-          ..fileOffset = fileOffset;
+    IfCaseMapEntry result = new IfCaseMapEntry(
+        prelude: prelude,
+        expression: expression,
+        patternGuard: patternGuard,
+        then: thenEntry,
+        otherwise: otherwiseEntry)
+      ..fileOffset = fileOffset;
     onConvertElement(this, result);
     return result;
   }
@@ -4880,10 +4890,14 @@ class IfCaseMapEntry extends TreeNode
   PatternGuard patternGuard;
   MapLiteralEntry then;
   MapLiteralEntry? otherwise;
-  late List<Statement> replacement;
+  List<Statement> prelude;
 
   IfCaseMapEntry(
-      this.expression, this.patternGuard, this.then, this.otherwise) {
+      {required this.prelude,
+      required this.expression,
+      required this.patternGuard,
+      required this.then,
+      this.otherwise}) {
     expression.parent = this;
     patternGuard.parent = this;
     then.parent = this;

pkg/front_end/lib/src/fasta/type_inference/inference_visitor.dart

@@ -2273,7 +2273,7 @@ class InferenceVisitorImpl extends InferenceVisitorBase
             createAndExpression(condition, guard, fileOffset: guard.fileOffset);
       }
       element.expression = condition;
-      element.replacement = [
+      element.prelude = [
         ...element.patternGuard.pattern.declaredVariables,
         ...matchingCache.declarations,
       ];
@@ -2829,7 +2829,7 @@ class InferenceVisitorImpl extends InferenceVisitorBase
     IfStatement ifStatement =
         _createIf(element.fileOffset, element.expression, thenBody, elseBody);
     libraryBuilder.loader.dataForTesting?.registerAlias(element, ifStatement);
-    body.addAll(element.replacement);
+    body.addAll(element.prelude);
     body.add(ifStatement);
   }
 
@@ -3165,7 +3165,7 @@ class InferenceVisitorImpl extends InferenceVisitorBase
     IfStatement ifStatement = _createIf(
         entry.fileOffset, entry.expression, thenStatement, elseStatement);
     libraryBuilder.loader.dataForTesting?.registerAlias(entry, ifStatement);
-    body.addAll(entry.replacement);
+    body.addAll(entry.prelude);
     body.add(ifStatement);
   }
 
@@ -4181,7 +4181,7 @@ class InferenceVisitorImpl extends InferenceVisitorBase
             createAndExpression(condition, guard, fileOffset: guard.fileOffset);
       }
       entry.expression = condition;
-      entry.replacement = [
+      entry.prelude = [
         ...entry.patternGuard.pattern.declaredVariables,
         ...matchingCache.declarations,
       ];

pkg/front_end/testcases/patterns/collections_entry_to_element_conversion_in_inference.dart

@@ -0,0 +1,59 @@
+// Copyright (c) 2023, the Dart project authors.  Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+Iterable<dynamic> test1(dynamic x, dynamic another) {
+  return {for (var [int i, int n] = x; i < n; i++) ...another};
+}
+
+Iterable<dynamic> test2(dynamic x, dynamic another) {
+  return {if (x case String _) ...another};
+}
+
+Iterable<dynamic> test3(dynamic x, dynamic another) {
+  return {for (var [int _] in x) ...another};
+}
+
+main() {
+  expectSetEquals(
+    test1([0, 2], {1, 2, 3}) as Set,
+    {1, 2, 3},
+  );
+  expectSetEquals(
+    test2([0, 0], {1, 2, 3}) as Set,
+    {},
+  );
+  expectThrows(() => test1([], {}));
+
+  expectSetEquals(
+    test2("foo", {1, 2, 3}) as Set,
+    {1, 2, 3},
+  );
+  expectSetEquals(
+    test2(false, {1, 2, 3}) as Set,
+    {},
+  );
+
+  expectSetEquals(
+    test3([[0], [1]], {1, 2, 3}) as Set,
+    {1, 2, 3},
+  );
+  expectThrows(() => test3([null], {}));
+}
+
+expectSetEquals(Set x, Set y) {
+  if (!x.containsAll(y) || !y.containsAll(x)) {
+    throw "Expected sets '${x}' and '${y}' to be equal.";
+  }
+}
+
+expectThrows(void Function() f) {
+  bool hasThrown = true;
+  try {
+    f();
+    hasThrown = false;
+  } catch (e) {}
+  if (!hasThrown) {
+    throw "Expected the function to throw.";
+  }
+}