[cfe] Support set/map disambiguation in pattern-for elements

Part of #49749

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

Author: chloestefantsova

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

@@ -4188,28 +4188,31 @@ class BodyBuilder extends StackListenerImpl
       assert(conditionStatement is EmptyStatement);
     }
     if (entry is MapLiteralEntry) {
-      TreeNode result;
-      ForMapEntry forMapEntry = result = forest.createForMapEntry(
-          offsetForToken(forToken), variables, condition, updates, entry);
+      ForMapEntry result;
       if (variableOrExpression is PatternVariableDeclaration) {
-        result = forest.createPatternForMapEntry(
-            offsetForToken(forToken), variableOrExpression, forMapEntry);
+        result = forest.createPatternForMapEntry(offsetForToken(forToken),
+            variableOrExpression, variables, condition, updates, entry);
+      } else {
+        result = forest.createForMapEntry(
+            offsetForToken(forToken), variables, condition, updates, entry);
       }
-      typeInferrer.assignedVariables.endNode(forMapEntry);
+      typeInferrer.assignedVariables.endNode(result);
       push(result);
     } else {
-      TreeNode result;
-      ForElement forElement = result = forest.createForElement(
-          offsetForToken(forToken),
-          variables,
-          condition,
-          updates,
-          toValue(entry));
+      ForElement result;
       if (variableOrExpression is PatternVariableDeclaration) {
         result = forest.createPatternForElement(
-            offsetForToken(forToken), variableOrExpression, forElement);
+            offsetForToken(forToken),
+            variableOrExpression,
+            variables,
+            condition,
+            updates,
+            toValue(entry));
+      } else {
+        result = forest.createForElement(offsetForToken(forToken), variables,
+            condition, updates, toValue(entry));
       }
-      typeInferrer.assignedVariables.endNode(forElement);
+      typeInferrer.assignedVariables.endNode(result);
       push(result);
     }
   }

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

@@ -299,8 +299,23 @@ class ForElement extends Expression with ControlFlowElement {
   }
 
   @override
-  void toTextInternal(AstPrinter state) {
-    // TODO(johnniwinther): Implement this.
+  void toTextInternal(AstPrinter printer) {
+    printer.write('for (');
+    for (int index = 0; index < variables.length; index++) {
+      if (index > 0) {
+        printer.write(', ');
+      }
+      printer.writeVariableDeclaration(variables[index],
+          includeModifiersAndType: index == 0);
+    }
+    printer.write('; ');
+    if (condition != null) {
+      printer.writeExpression(condition!);
+    }
+    printer.write('; ');
+    printer.writeExpressions(updates);
+    printer.write(') ');
+    printer.writeExpression(body);
   }
 }
 
@@ -658,8 +673,23 @@ class ForMapEntry extends TreeNode with ControlFlowMapEntry {
   }
 
   @override
-  void toTextInternal(AstPrinter state) {
-    // TODO(johnniwinther): Implement this.
+  void toTextInternal(AstPrinter printer) {
+    printer.write('for (');
+    for (int index = 0; index < variables.length; index++) {
+      if (index > 0) {
+        printer.write(', ');
+      }
+      printer.writeVariableDeclaration(variables[index],
+          includeModifiersAndType: index == 0);
+    }
+    printer.write('; ');
+    if (condition != null) {
+      printer.writeExpression(condition!);
+    }
+    printer.write('; ');
+    printer.writeExpressions(updates);
+    printer.write(') ');
+    body.toTextInternal(printer);
   }
 }
 
@@ -859,6 +889,9 @@ bool isConvertibleToMapEntry(Expression element) {
   if (element is ForElement) {
     return isConvertibleToMapEntry(element.body);
   }
+  if (element is PatternForElement) {
+    return isConvertibleToMapEntry(element.body);
+  }
   if (element is ForInElement) {
     return isConvertibleToMapEntry(element.body);
   }
@@ -903,6 +936,17 @@ MapLiteralEntry convertToMapEntry(Expression element, InferenceHelper helper,
     return result;
   }
   if (element is ForElement) {
+    if (element is PatternForElement) {
+      PatternForMapEntry result = new PatternForMapEntry(
+          element.patternVariableDeclaration,
+          element.variables,
+          element.condition,
+          element.updates,
+          convertToMapEntry(element.body, helper, onConvertElement))
+        ..fileOffset = element.fileOffset;
+      onConvertElement(element, result);
+      return result;
+    }
     ForMapEntry result = new ForMapEntry(
         element.variables,
         element.condition,

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

@@ -313,10 +313,14 @@ class Forest {
   PatternForElement createPatternForElement(
       int fileOffset,
       PatternVariableDeclaration patternVariableDeclaration,
-      ForElement forElement) {
+      List<VariableDeclaration> variables,
+      Expression? condition,
+      List<Expression> updates,
+      Expression body) {
     // ignore: unnecessary_null_comparison
     assert(fileOffset != null);
-    return new PatternForElement(patternVariableDeclaration, forElement)
+    return new PatternForElement(
+        patternVariableDeclaration, variables, condition, updates, body)
       ..fileOffset = fileOffset;
   }
 
@@ -335,10 +339,14 @@ class Forest {
   PatternForMapEntry createPatternForMapEntry(
       int fileOffset,
       PatternVariableDeclaration patternVariableDeclaration,
-      ForMapEntry forMapEntry) {
+      List<VariableDeclaration> variables,
+      Expression? condition,
+      List<Expression> updates,
+      MapLiteralEntry body) {
     // ignore: unnecessary_null_comparison
     assert(fileOffset != null);
-    return new PatternForMapEntry(patternVariableDeclaration, forMapEntry)
+    return new PatternForMapEntry(
+        patternVariableDeclaration, variables, condition, updates, body)
       ..fileOffset = fileOffset;
   }
 

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

@@ -4915,12 +4915,26 @@ class IfCaseMapEntry extends TreeNode
   }
 }
 
-class PatternForElement extends InternalExpression with ControlFlowElement {
+class PatternForElement extends InternalExpression
+    with ControlFlowElement
+    implements ForElement {
   PatternVariableDeclaration patternVariableDeclaration;
-  ForElement forElement;
   late List<Statement> replacement;
 
-  PatternForElement(this.patternVariableDeclaration, this.forElement);
+  @override
+  final List<VariableDeclaration> variables; // May be empty, but not null.
+
+  @override
+  Expression? condition; // May be null.
+
+  @override
+  final List<Expression> updates; // May be empty, but not null.
+
+  @override
+  Expression body;
+
+  PatternForElement(this.patternVariableDeclaration, this.variables,
+      this.condition, this.updates, this.body);
 
   @override
   ExpressionInferenceResult acceptInference(
@@ -4930,22 +4944,23 @@ class PatternForElement extends InternalExpression with ControlFlowElement {
 
   @override
   void toTextInternal(AstPrinter printer) {
+    patternVariableDeclaration.toTextInternal(printer);
     printer.write('for (');
-    for (int index = 0; index < forElement.variables.length; index++) {
+    for (int index = 0; index < variables.length; index++) {
       if (index > 0) {
         printer.write(', ');
       }
-      printer.writeVariableDeclaration(forElement.variables[index],
+      printer.writeVariableDeclaration(variables[index],
           includeModifiersAndType: index == 0);
     }
     printer.write('; ');
-    if (forElement.condition != null) {
-      printer.writeExpression(forElement.condition!);
+    if (condition != null) {
+      printer.writeExpression(condition!);
     }
     printer.write('; ');
-    printer.writeExpressions(forElement.updates);
+    printer.writeExpressions(updates);
     printer.write(') ');
-    printer.writeExpression(forElement.body);
+    printer.writeExpression(body);
   }
 
   @override
@@ -4961,12 +4976,25 @@ class PatternForElement extends InternalExpression with ControlFlowElement {
 }
 
 class PatternForMapEntry extends TreeNode
-    with InternalTreeNode, ControlFlowMapEntry {
+    with InternalTreeNode, ControlFlowMapEntry
+    implements ForMapEntry {
   PatternVariableDeclaration patternVariableDeclaration;
-  ForMapEntry forMapEntry;
   late List<Statement> replacement;
 
-  PatternForMapEntry(this.patternVariableDeclaration, this.forMapEntry);
+  @override
+  final List<VariableDeclaration> variables;
+
+  @override
+  Expression? condition;
+
+  @override
+  final List<Expression> updates;
+
+  @override
+  MapLiteralEntry body;
+
+  PatternForMapEntry(this.patternVariableDeclaration, this.variables,
+      this.condition, this.updates, this.body);
 
   ExpressionInferenceResult acceptInference(
       InferenceVisitorImpl visitor, DartType typeContext) {
@@ -4975,24 +5003,23 @@ class PatternForMapEntry extends TreeNode
 
   @override
   void toTextInternal(AstPrinter printer) {
+    patternVariableDeclaration.toTextInternal(printer);
     printer.write('for (');
-    for (int index = 0; index < forMapEntry.variables.length; index++) {
+    for (int index = 0; index < variables.length; index++) {
       if (index > 0) {
         printer.write(', ');
       }
-      printer.writeVariableDeclaration(forMapEntry.variables[index],
+      printer.writeVariableDeclaration(variables[index],
           includeModifiersAndType: index == 0);
     }
     printer.write('; ');
-    if (forMapEntry.condition != null) {
-      printer.writeExpression(forMapEntry.condition!);
+    if (condition != null) {
+      printer.writeExpression(condition!);
     }
     printer.write('; ');
-    printer.writeExpressions(forMapEntry.updates);
+    printer.writeExpressions(updates);
     printer.write(') ');
-    printer.writeExpression(forMapEntry.body.key);
-    printer.write(': ');
-    printer.writeExpression(forMapEntry.body.value);
+    body.toTextInternal(printer);
   }
 
   @override