[cfe] Add basic support for if-case elements in maps

Part of #49749

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

Author: chloestefantsova

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

@@ -6597,17 +6597,28 @@ class BodyBuilder extends StackListenerImpl
     exitLocalScope(expectedScopeKinds: const [ScopeKind.ifElement]);
     Token ifToken = pop() as Token;
 
+    PatternGuard? patternGuard = condition.patternGuard;
     TreeNode node;
     if (thenEntry is MapLiteralEntry) {
       if (elseEntry is MapLiteralEntry) {
-        node = forest.createIfMapEntry(offsetForToken(ifToken),
-            condition.expression, thenEntry, elseEntry);
+        if (patternGuard == null) {
+          node = forest.createIfMapEntry(offsetForToken(ifToken),
+              condition.expression, thenEntry, elseEntry);
+        } else {
+          node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
+              condition.expression, patternGuard, thenEntry, elseEntry);
+        }
       } else if (elseEntry is ControlFlowElement) {
         MapLiteralEntry? elseMapEntry = elseEntry
             .toMapLiteralEntry(typeInferrer.assignedVariables.reassignInfo);
         if (elseMapEntry != null) {
-          node = forest.createIfMapEntry(offsetForToken(ifToken),
-              condition.expression, thenEntry, elseMapEntry);
+          if (patternGuard == null) {
+            node = forest.createIfMapEntry(offsetForToken(ifToken),
+                condition.expression, thenEntry, elseMapEntry);
+          } else {
+            node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
+                condition.expression, patternGuard, thenEntry, elseMapEntry);
+          }
         } else {
           int offset = elseEntry.fileOffset;
           node = new MapLiteralEntry(
@@ -6631,8 +6642,13 @@ class BodyBuilder extends StackListenerImpl
         MapLiteralEntry? thenMapEntry = thenEntry
             .toMapLiteralEntry(typeInferrer.assignedVariables.reassignInfo);
         if (thenMapEntry != null) {
-          node = forest.createIfMapEntry(offsetForToken(ifToken),
-              condition.expression, thenMapEntry, elseEntry);
+          if (patternGuard == null) {
+            node = forest.createIfMapEntry(offsetForToken(ifToken),
+                condition.expression, thenMapEntry, elseEntry);
+          } else {
+            node = forest.createIfCaseMapEntry(offsetForToken(ifToken),
+                condition.expression, patternGuard, thenMapEntry, elseEntry);
+          }
         } else {
           int offset = thenEntry.fileOffset;
           node = new MapLiteralEntry(

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

@@ -4861,6 +4861,7 @@ class IfCaseMapEntry extends TreeNode
   PatternGuard patternGuard;
   MapLiteralEntry then;
   MapLiteralEntry? otherwise;
+  late List<Statement> replacement;
 
   IfCaseMapEntry(
       this.expression, this.patternGuard, this.then, this.otherwise) {

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

@@ -2197,9 +2197,9 @@ class InferenceVisitorImpl extends InferenceVisitorBase
       int? stackBase;
       assert(checkStackBase(element, stackBase = stackHeight));
 
-      CollectionElementInferenceContext context =
-          new CollectionElementInferenceContext(
-              typeContext: inferredTypeArgument,
+      ListAndSetElementInferenceContext context =
+          new ListAndSetElementInferenceContext(
+              inferredTypeArgument: inferredTypeArgument,
               inferredSpreadTypes: inferredSpreadTypes,
               inferredConditionTypes: inferredConditionTypes);
       analyzeIfCaseElement(
@@ -2972,6 +2972,9 @@ class InferenceVisitorImpl extends InferenceVisitorBase
           entry, receiverType, keyType, valueType, result, body);
     } else if (entry is IfMapEntry) {
       _translateIfEntry(entry, receiverType, keyType, valueType, result, body);
+    } else if (entry is IfCaseMapEntry) {
+      _translateIfCaseEntry(
+          entry, receiverType, keyType, valueType, result, body);
     } else if (entry is ForMapEntry) {
       _translateForEntry(entry, receiverType, keyType, valueType, result, body);
     } else if (entry is ForInMapEntry) {
@@ -3016,6 +3019,35 @@ class InferenceVisitorImpl extends InferenceVisitorBase
     body.add(ifStatement);
   }
 
+  void _translateIfCaseEntry(
+      IfCaseMapEntry entry,
+      InterfaceType receiverType,
+      DartType keyType,
+      DartType valueType,
+      VariableDeclaration result,
+      List<Statement> body) {
+    List<Statement> thenBody = [];
+    _translateEntry(
+        entry.then, receiverType, keyType, valueType, result, thenBody);
+    List<Statement>? elseBody;
+    if (entry.otherwise != null) {
+      _translateEntry(entry.otherwise!, receiverType, keyType, valueType,
+          result, elseBody = <Statement>[]);
+    }
+    Statement thenStatement =
+        thenBody.length == 1 ? thenBody.first : _createBlock(thenBody);
+    Statement? elseStatement;
+    if (elseBody != null && elseBody.isNotEmpty) {
+      elseStatement =
+          elseBody.length == 1 ? elseBody.first : _createBlock(elseBody);
+    }
+    IfStatement ifStatement = _createIf(
+        entry.fileOffset, entry.expression, thenStatement, elseStatement);
+    libraryBuilder.loader.dataForTesting?.registerAlias(entry, ifStatement);
+    body.addAll(entry.replacement);
+    body.add(ifStatement);
+  }
+
   void _translateForEntry(
       ForMapEntry entry,
       InterfaceType receiverType,
@@ -3911,7 +3943,18 @@ class InferenceVisitorImpl extends InferenceVisitorBase
       flowAnalysis.ifStatement_end(entry.otherwise != null);
       return entry;
     } else if (entry is IfCaseMapEntry) {
-      // TODO(cstefantsova): Pass an appropriate context message.
+      int? stackBase;
+      assert(checkStackBase(entry, stackBase = stackHeight));
+
+      MapEntryInferenceContext context = new MapEntryInferenceContext(
+          inferredKeyType: inferredKeyType,
+          inferredValueType: inferredValueType,
+          spreadContext: spreadContext,
+          actualTypes: actualTypes,
+          actualTypesForSet: actualTypesForSet,
+          offsets: offsets,
+          inferredSpreadTypes: inferredSpreadTypes,
+          inferredConditionTypes: inferredConditionTypes);
       analyzeIfCaseElement(
           node: entry,
           expression: entry.expression,
@@ -3924,9 +3967,86 @@ class InferenceVisitorImpl extends InferenceVisitorBase
           guard: entry.patternGuard.guard,
           ifTrue: entry.then,
           ifFalse: entry.otherwise,
-          context: null);
-      // TODO(cstefantsova): Implement inference for if-case map entries.
-      throw new UnimplementedError();
+          context: context);
+      if (entry.otherwise != null) {
+        DartType actualValueType = actualTypes.removeLast();
+        DartType actualKeyType = actualTypes.removeLast();
+        int length = actualTypes.length;
+        actualTypes[length - 2] = typeSchemaEnvironment.getStandardUpperBound(
+            actualKeyType, actualTypes[length - 2],
+            isNonNullableByDefault: libraryBuilder.isNonNullableByDefault);
+        actualTypes[length - 1] = typeSchemaEnvironment.getStandardUpperBound(
+            actualValueType, actualTypes[length - 1],
+            isNonNullableByDefault: libraryBuilder.isNonNullableByDefault);
+        DartType actualTypeForSet = actualTypesForSet.removeLast();
+        int lengthForSet = actualTypesForSet.length;
+        actualTypesForSet[lengthForSet - 1] =
+            typeSchemaEnvironment.getStandardUpperBound(
+                actualTypeForSet, actualTypesForSet[lengthForSet - 1],
+                isNonNullableByDefault: libraryBuilder.isNonNullableByDefault);
+      }
+
+      assert(checkStack(entry, stackBase, [
+        /* ifFalse = */ unionOfKinds(
+            [ValueKinds.MapLiteralEntryOrNull, ValueKinds.ExpressionOrNull]),
+        /* ifTrue = */ unionOfKinds(
+            [ValueKinds.MapLiteralEntry, ValueKinds.Expression]),
+        /* guard = */ ValueKinds.ExpressionOrNull,
+        /* pattern = */ ValueKinds.Pattern,
+        /* scrutinee = */ ValueKinds.Expression,
+      ]));
+
+      Object? rewrite = popRewrite(NullValues.Expression);
+      if (!identical(entry.otherwise, rewrite)) {
+        entry.otherwise = (rewrite as MapLiteralEntry?)?..parent = entry;
+      }
+
+      rewrite = popRewrite();
+      if (!identical(entry.then, rewrite)) {
+        entry.then = (rewrite as MapLiteralEntry)..parent = entry;
+      }
+
+      PatternGuard patternGuard = entry.patternGuard;
+      rewrite = popRewrite(NullValues.Expression);
+      if (!identical(patternGuard.guard, rewrite)) {
+        patternGuard.guard = (rewrite as Expression?)?..parent = patternGuard;
+      }
+
+      rewrite = popRewrite();
+      if (!identical(patternGuard.pattern, rewrite)) {
+        patternGuard.pattern = (rewrite as Pattern)..parent = patternGuard;
+      }
+
+      rewrite = popRewrite();
+      if (!identical(entry.expression, rewrite)) {
+        entry.expression = (rewrite as Expression)..parent = patternGuard;
+      }
+
+      MatchingCache matchingCache = createMatchingCache();
+      MatchingExpressionVisitor matchingExpressionVisitor =
+          new MatchingExpressionVisitor(matchingCache);
+      // TODO(cstefantsova): Provide a more precise scrutinee type.
+      CacheableExpression matchedExpression = matchingCache
+          .createRootExpression(entry.expression, const DynamicType());
+      DelayedExpression matchingExpression = matchingExpressionVisitor
+          .visitPattern(entry.patternGuard.pattern, matchedExpression);
+
+      matchingExpression.registerUse();
+
+      Expression condition =
+          matchingExpression.createExpression(typeSchemaEnvironment);
+      Expression? guard = entry.patternGuard.guard;
+      if (guard != null) {
+        condition =
+            createAndExpression(condition, guard, fileOffset: guard.fileOffset);
+      }
+      entry.expression = condition;
+      entry.replacement = [
+        ...entry.patternGuard.pattern.declaredVariables,
+        ...matchingCache.declarations,
+      ];
+
+      return entry;
     } else if (entry is ForMapEntry) {
       // TODO(johnniwinther): Use _visitStatements instead.
       List<VariableDeclaration>? variables;
@@ -10587,17 +10707,37 @@ class InferenceVisitorImpl extends InferenceVisitorBase
   }
 
   @override
-  void dispatchCollectionElement(covariant Expression element,
+  void dispatchCollectionElement(covariant TreeNode element,
       covariant CollectionElementInferenceContext context) {
-    ExpressionInferenceResult inferenceResult = inferElement(
-        element,
-        context.typeContext,
-        context.inferredSpreadTypes,
-        context.inferredConditionTypes);
-    // TODO(cstefantsova): Should the key to the map be [element] instead?
-    context.inferredConditionTypes[inferenceResult.expression] =
-        inferenceResult.inferredType;
-    pushRewrite(inferenceResult.expression);
+    if (element is Expression) {
+      context as ListAndSetElementInferenceContext;
+      ExpressionInferenceResult inferenceResult = inferElement(
+          element,
+          context.inferredTypeArgument,
+          context.inferredSpreadTypes,
+          context.inferredConditionTypes);
+      // TODO(cstefantsova): Should the key to the map be [element] instead?
+      context.inferredConditionTypes[inferenceResult.expression] =
+          inferenceResult.inferredType;
+      pushRewrite(inferenceResult.expression);
+    } else if (element is MapLiteralEntry) {
+      context as MapEntryInferenceContext;
+      element = inferMapEntry(
+          element,
+          element.parent!,
+          context.inferredKeyType,
+          context.inferredValueType,
+          context.spreadContext,
+          context.actualTypes,
+          context.actualTypesForSet,
+          context.inferredSpreadTypes,
+          context.inferredConditionTypes,
+          context.offsets);
+      pushRewrite(element);
+    } else {
+      problems.unsupported(
+          "${element.runtimeType}", element.fileOffset, helper.uri);
+    }
   }
 
   @override
@@ -10811,13 +10951,46 @@ extension on SwitchCase {
   }
 }
 
-class CollectionElementInferenceContext {
-  DartType typeContext;
+abstract class CollectionElementInferenceContext {
   Map<TreeNode, DartType> inferredSpreadTypes;
   Map<Expression, DartType> inferredConditionTypes;
 
   CollectionElementInferenceContext(
-      {required this.typeContext,
-      required this.inferredSpreadTypes,
+      {required this.inferredSpreadTypes,
       required this.inferredConditionTypes});
 }
+
+class ListAndSetElementInferenceContext
+    extends CollectionElementInferenceContext {
+  DartType inferredTypeArgument;
+
+  ListAndSetElementInferenceContext(
+      {required this.inferredTypeArgument,
+      required Map<TreeNode, DartType> inferredSpreadTypes,
+      required Map<Expression, DartType> inferredConditionTypes})
+      : super(
+            inferredSpreadTypes: inferredSpreadTypes,
+            inferredConditionTypes: inferredConditionTypes);
+}
+
+class MapEntryInferenceContext extends CollectionElementInferenceContext {
+  DartType inferredKeyType;
+  DartType inferredValueType;
+  DartType spreadContext;
+  List<DartType> actualTypes;
+  List<DartType> actualTypesForSet;
+  _MapLiteralEntryOffsets offsets;
+
+  MapEntryInferenceContext(
+      {required this.inferredKeyType,
+      required this.inferredValueType,
+      required this.spreadContext,
+      required this.actualTypes,
+      required this.actualTypesForSet,
+      required this.offsets,
+      required Map<TreeNode, DartType> inferredSpreadTypes,
+      required Map<Expression, DartType> inferredConditionTypes})
+      : super(
+            inferredSpreadTypes: inferredSpreadTypes,
+            inferredConditionTypes: inferredConditionTypes);
+}

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

@@ -10,6 +10,8 @@ import '../kernel/internal_ast.dart' as type;
 class NullValues {
   static const NullValue<type.Expression> Expression =
       const NullValue<type.Expression>();
+  static const NullValue<type.MapLiteralEntry> MapLiteralEntry =
+      const NullValue<type.MapLiteralEntry>();
   static const NullValue<type.Pattern> Pattern =
       const NullValue<type.Pattern>();
   static const NullValue<type.Statement> Statement =
@@ -23,6 +25,10 @@ class ValueKinds {
       const SingleValueKind<type.Expression>();
   static const SingleValueKind<type.Expression> ExpressionOrNull =
       const SingleValueKind<type.Expression>(NullValues.Expression);
+  static const SingleValueKind<type.MapLiteralEntry> MapLiteralEntry =
+      const SingleValueKind<type.MapLiteralEntry>();
+  static const SingleValueKind<type.MapLiteralEntry> MapLiteralEntryOrNull =
+      const SingleValueKind<type.MapLiteralEntry>(NullValues.MapLiteralEntry);
   static const SingleValueKind<type.MapPatternEntry> MapPatternEntry =
       const SingleValueKind<type.MapPatternEntry>();
   static const SingleValueKind<type.Pattern> Pattern =

pkg/front_end/testcases/patterns/simple_if_case_map_entries.dart

@@ -0,0 +1,53 @@
+// 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.
+
+test1(dynamic x) => {1: 1, if (x case int y) 2: y, 3: 3};
+
+test2(dynamic x) => {1: 1, if (x case String y) 2: y else 2: null, 3: 3};
+
+test3(dynamic x) => {1: 1, if (x case bool b when b) 2: b, 3: 3};
+
+main() {
+  expectEquals(
+    mapToString(test1(0)),
+    mapToString({1: 1, 2: 0, 3: 3}),
+  );
+  expectEquals(
+    mapToString(test1("foo")),
+    mapToString({1: 1, 3: 3}),
+  );
+
+  expectEquals(
+    mapToString(test2("foo")),
+    mapToString({1: 1, 2: "foo", 3: 3}),
+  );
+  expectEquals(
+    mapToString(test2(false)),
+    mapToString({1: 1, 2: null, 3: 3}),
+  );
+
+  expectEquals(
+    mapToString(test3(true)),
+    mapToString({1: 1, 2: true, 3: 3}),
+  );
+  expectEquals(
+    mapToString(test3(false)),
+    mapToString({1: 1, 3: 3}),
+  );
+}
+
+expectEquals(x, y) {
+  if (x != y) {
+    throw "Expected '${x}' to be equal to '${y}'.";
+  }
+}
+
+mapToString(Map<dynamic, dynamic> map) {
+  List<String> entryStrings = [
+    for (var entry in map.entries)
+      "${entry.key}:${entry.value}"
+  ];
+  entryStrings.sort();
+  return "{${entryStrings.join(',')}}";
+}