[analyzer][cfe] Implement flow analysis for null-aware map entries

Closes #56786

Change-Id: I738c98b6f4e632cfbbe51221bbc3547edbc718fe
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/386800
Commit-Queue: Chloe Stefantsova <[email protected]>
Reviewed-by: Paul Berry <[email protected]>

Author: chloestefantsova

pkg/_fe_analyzer_shared/lib/src/flow_analysis/flow_analysis.dart

@@ -703,6 +703,13 @@ abstract class FlowAnalysis<Node extends Object, Statement extends Node,
   /// not be called until after processing the method call to `z(x)`.
   void nullAwareAccess_rightBegin(Expression? target, Type targetType);
 
+  /// Call this method after visiting the value of a null-aware map entry.
+  void nullAwareMapEntry_end({required bool isKeyNullAware});
+
+  /// Call this method after visiting the key of a null-aware map entry.
+  void nullAwareMapEntry_valueBegin(Expression key, Type keyType,
+      {required bool isKeyNullAware});
+
   /// Call this method before visiting the subpattern of a null-check or a
   /// null-assert pattern. [isAssert] indicates whether the pattern is a
   /// null-check or a null-assert pattern.
@@ -1655,6 +1662,22 @@ class FlowAnalysisDebug<Node extends Object, Statement extends Node,
         () => _wrapped.nullAwareAccess_rightBegin(target, targetType));
   }
 
+  @override
+  void nullAwareMapEntry_end({required bool isKeyNullAware}) {
+    return _wrap('nullAwareMapEntry_end(isKeyNullAware: $isKeyNullAware)',
+        () => _wrapped.nullAwareMapEntry_end(isKeyNullAware: isKeyNullAware));
+  }
+
+  @override
+  void nullAwareMapEntry_valueBegin(Expression key, Type keyType,
+      {required bool isKeyNullAware}) {
+    _wrap(
+        'nullAwareMapEntry_valueBegin($key, $keyType, '
+        'isKeyNullAware: $isKeyNullAware)',
+        () => _wrapped.nullAwareMapEntry_valueBegin(key, keyType,
+            isKeyNullAware: isKeyNullAware));
+  }
+
   @override
   bool nullCheckOrAssertPattern_begin(
       {required bool isAssert, required Type matchedValueType}) {
@@ -5032,6 +5055,35 @@ class _FlowAnalysisImpl<Node extends Object, Statement extends Node,
     }
   }
 
+  @override
+  void nullAwareMapEntry_end({required bool isKeyNullAware}) {
+    if (!isKeyNullAware) return;
+    _NullAwareMapEntryContext<Type> context =
+        _stack.removeLast() as _NullAwareMapEntryContext<Type>;
+    _current = _join(_current, context._shortcutState).unsplit();
+  }
+
+  @override
+  void nullAwareMapEntry_valueBegin(Expression key, Type keyType,
+      {required bool isKeyNullAware}) {
+    if (!isKeyNullAware) return;
+    _Reference<Type>? keyReference = _getExpressionReference(key);
+    FlowModel<Type> shortcutState;
+    _current = _current.split();
+    if (keyReference != null) {
+      ExpressionInfo<Type> expressionInfo =
+          _current.tryMarkNonNullable(this, keyReference);
+      _current = expressionInfo.ifTrue;
+      shortcutState = expressionInfo.ifFalse;
+    } else {
+      shortcutState = _current;
+    }
+    if (operations.classifyType(keyType) == TypeClassification.nonNullable) {
+      shortcutState = shortcutState.setUnreachable();
+    }
+    _stack.add(new _NullAwareMapEntryContext<Type>(shortcutState));
+  }
+
   @override
   bool nullCheckOrAssertPattern_begin(
       {required bool isAssert, required Type matchedValueType}) {
@@ -6640,6 +6692,13 @@ class _LegacyTypePromotion<Node extends Object, Statement extends Node,
   @override
   void nullAwareAccess_rightBegin(Expression? target, Type targetType) {}
 
+  @override
+  void nullAwareMapEntry_end({required bool isKeyNullAware}) {}
+
+  @override
+  void nullAwareMapEntry_valueBegin(Expression key, Type keyType,
+      {required bool isKeyNullAware}) {}
+
   @override
   bool nullCheckOrAssertPattern_begin(
           {required bool isAssert, required Type matchedValueType}) =>
@@ -6922,6 +6981,25 @@ class _NullAwareAccessContext<Type extends Object>
   String get _debugType => '_NullAwareAccessContext';
 }
 
+/// [_FlowContext] representing a null-aware map entry (`{?a: ?b}`).
+///
+/// This context should only be created for a null-aware map entry that has a
+/// null-aware key.
+class _NullAwareMapEntryContext<Type extends Object> extends _FlowContext {
+  /// The state if the operation short-cuts (i.e. if the key expression was
+  /// `null`.
+  final FlowModel<Type> _shortcutState;
+
+  _NullAwareMapEntryContext(this._shortcutState);
+
+  @override
+  Map<String, Object?> get _debugFields =>
+      super._debugFields..['shortcutState'] = _shortcutState;
+
+  @override
+  String get _debugType => '_NullAwareMapEntryContext';
+}
+
 /// Specialization of [ExpressionInfo] for the case where the expression is a
 /// `null` literal.
 class _NullInfo<Type extends Object> extends ExpressionInfo<Type> {

pkg/_fe_analyzer_shared/test/flow_analysis/flow_analysis_test.dart

@@ -8606,6 +8606,82 @@ main() {
       });
     });
 
+    group('Null-aware map entry:', () {
+      test('Promotes key within value', () {
+        var a = Var('a');
+
+        h.run([
+          declare(a, type: 'String?', initializer: expr('String?')),
+          mapLiteral(keyType: 'String', valueType: 'dynamic', [
+            mapEntry(a, checkPromoted(a, 'String'), isKeyNullAware: true),
+          ]),
+          checkNotPromoted(a),
+        ]);
+      });
+
+      test('Non-null-aware key', () {
+        var a = Var('a');
+
+        h.run([
+          declare(a, type: 'String?', initializer: expr('String?')),
+          mapLiteral(keyType: 'String?', valueType: 'dynamic', [
+            mapEntry(a, checkNotPromoted(a), isKeyNullAware: false),
+          ]),
+          checkNotPromoted(a),
+        ]);
+      });
+
+      test('Promotes', () {
+        var a = Var('a');
+        var x = Var('x');
+
+        h.run([
+          declare(a, type: 'String', initializer: expr('String')),
+          declare(x, type: 'num', initializer: expr('num')),
+          mapLiteral(keyType: 'String', valueType: 'dynamic', [
+            mapEntry(a, x.as_('int'), isKeyNullAware: true),
+          ]),
+          checkPromoted(x, 'int'),
+        ]);
+      });
+
+      test('Affects promotion', () {
+        var a = Var('a');
+        var x = Var('x');
+
+        h.run([
+          declare(a, type: 'String?', initializer: expr('String?')),
+          declare(x, type: 'num', initializer: expr('num')),
+          mapLiteral(keyType: 'String', valueType: 'dynamic', [
+            mapEntry(a, x.as_('int'), isKeyNullAware: true),
+          ]),
+          checkNotPromoted(x),
+        ]);
+      });
+
+      test('Unreachable', () {
+        var a = Var('a');
+        h.run([
+          declare(a, type: 'String', initializer: expr('String')),
+          mapLiteral(keyType: 'String', valueType: 'dynamic', [
+            mapEntry(a, throw_(expr('Object')), isKeyNullAware: true),
+          ]),
+          checkReachable(false),
+        ]);
+      });
+
+      test('Reachable', () {
+        var a = Var('a');
+        h.run([
+          declare(a, type: 'String?', initializer: expr('String?')),
+          mapLiteral(keyType: 'String', valueType: 'dynamic', [
+            mapEntry(a, throw_(expr('Object')), isKeyNullAware: true),
+          ]),
+          checkReachable(true),
+        ]);
+      });
+    });
+
     group('Map pattern:', () {
       test('Promotes', () {
         var x = Var('x');

pkg/_fe_analyzer_shared/test/mini_ast.dart

@@ -299,11 +299,12 @@ Expression localFunction(List<ProtoStatement> body) {
 }
 
 /// Creates a map entry containing the given [key] and [value] subexpressions.
-CollectionElement mapEntry(ProtoExpression key, ProtoExpression value) {
+CollectionElement mapEntry(ProtoExpression key, ProtoExpression value,
+    {bool isKeyNullAware = false}) {
   var location = computeLocation();
   return MapEntry._(key.asExpression(location: location),
       value.asExpression(location: location),
-      location: location);
+      isKeyNullAware: isKeyNullAware, location: location);
 }
 
 /// Creates a map literal containing the given [elements].
@@ -2527,8 +2528,10 @@ abstract class LValue extends Expression {
 class MapEntry extends CollectionElement {
   final Expression key;
   final Expression value;
+  final bool isKeyNullAware;
 
-  MapEntry._(this.key, this.value, {required super.location});
+  MapEntry._(this.key, this.value,
+      {required this.isKeyNullAware, required super.location});
 
   @override
   void preVisit(PreVisitor visitor) {
@@ -2537,7 +2540,7 @@ class MapEntry extends CollectionElement {
   }
 
   @override
-  String toString() => '$key: $value';
+  String toString() => '${isKeyNullAware ? '?' : ''}$key: $value';
 
   @override
   void visit(Harness h, CollectionElementContext context) {
@@ -2550,8 +2553,11 @@ class MapEntry extends CollectionElement {
       default:
         keySchema = valueSchema = h.operations.unknownType;
     }
-    h.typeAnalyzer.analyzeExpression(key, keySchema);
+    var keyType = h.typeAnalyzer.analyzeExpression(key, keySchema);
+    h.flow.nullAwareMapEntry_valueBegin(key, keyType,
+        isKeyNullAware: isKeyNullAware);
     h.typeAnalyzer.analyzeExpression(value, valueSchema);
+    h.flow.nullAwareMapEntry_end(isKeyNullAware: isKeyNullAware);
     h.irBuilder.apply(
         'mapEntry', [Kind.expression, Kind.expression], Kind.collectionElement,
         location: location);

pkg/analyzer/lib/src/generated/resolver.dart

@@ -3128,14 +3128,18 @@ class ResolverVisitor extends ThrowingAstVisitor<void>
       {CollectionLiteralContext? context}) {
     inferenceLogWriter?.enterElement(node);
     checkUnreachableNode(node);
-    analyzeExpression(node.key,
+    var keyType = analyzeExpression(node.key,
         SharedTypeSchemaView(context?.keyType ?? UnknownInferredType.instance));
     popRewrite();
+    flowAnalysis.flow?.nullAwareMapEntry_valueBegin(node.key, keyType,
+        isKeyNullAware: node.keyQuestion != null);
     analyzeExpression(
         node.value,
         SharedTypeSchemaView(
             context?.valueType ?? UnknownInferredType.instance));
     popRewrite();
+    flowAnalysis.flow
+        ?.nullAwareMapEntry_end(isKeyNullAware: node.keyQuestion != null);
     inferenceLogWriter?.exitElement(node);
   }
 

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

@@ -4595,9 +4595,6 @@ class InferenceVisitorImpl extends InferenceVisitorBase
       Map<TreeNode, DartType> inferredSpreadTypes,
       Map<Expression, DartType> inferredConditionTypes,
       _MapLiteralEntryOffsets offsets) {
-    // TODO(cstefantsova): Make sure flow analysis is invoked here when it's
-    // implemented.
-
     DartType adjustedInferredKeyType = entry.isKeyNullAware
         ? inferredKeyType.withDeclaredNullability(Nullability.nullable)
         : inferredKeyType;
@@ -4610,6 +4607,10 @@ class InferenceVisitorImpl extends InferenceVisitorBase
         .expression;
     entry.key = key..parent = entry;
 
+    flowAnalysis.nullAwareMapEntry_valueBegin(
+        key, new SharedTypeView(keyInferenceResult.inferredType),
+        isKeyNullAware: entry.isKeyNullAware);
+
     DartType adjustedInferredValueType = entry.isValueNullAware
         ? inferredValueType.withDeclaredNullability(Nullability.nullable)
         : inferredValueType;
@@ -4627,6 +4628,8 @@ class InferenceVisitorImpl extends InferenceVisitorBase
 
     offsets.mapEntryOffset = entry.fileOffset;
 
+    flowAnalysis.nullAwareMapEntry_end(isKeyNullAware: entry.isKeyNullAware);
+
     return entry;
   }
 

pkg/front_end/testcases/null_aware_elements/flow_analysis_null_aware_map_entires.dart

@@ -0,0 +1,76 @@
+// Copyright (c) 2024, 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(String a, num x) {
+  <String, dynamic>{?a: (x as int)};
+  x..expectStaticType<Exactly<int>>();
+}
+
+test2(String? a, num x) {
+  <String, dynamic>{?a: (x as int)};
+  x..expectStaticType<Exactly<num>>();
+}
+
+test3(String a, bool b, num x) {
+  if (b) {
+    x as int;
+  } else {
+    <String, dynamic>{?a: (throw 0)};
+    // Unreachable.
+  }
+  x..expectStaticType<Exactly<int>>();
+}
+
+test4(String? a, bool b, num x) {
+  if (b) {
+    x as int;
+  } else {
+    <String, dynamic>{?a: (throw 0)};
+    // Reachable.
+  }
+  x..expectStaticType<Exactly<num>>();
+}
+
+test5(String? a) {
+  return {?a: a..expectStaticType<Exactly<String>>()};
+}
+
+test6(String? a) {
+  return {a: a..expectStaticType<Exactly<String?>>()};
+}
+
+extension E<X> on X {
+  void expectStaticType<Y extends Exactly<X>>() {}
+}
+
+typedef Exactly<X> = X Function(X);
+
+void expectThrows(void Function() f) {
+  bool hasThrown;
+  try {
+    f();
+    hasThrown = false;
+  } catch (e) {
+    hasThrown = true;
+  }
+
+  if (!hasThrown) {
+    throw "Expected the function to throw an exception.";
+  }
+}
+
+main() {
+  test1("", 0);
+
+  test2("", 0);
+  test2(null, 0);
+
+  test3("", true, 0);
+  expectThrows(() => test3("", false, 0));
+
+  test4("", true, 0);
+  expectThrows(() => test4("", false, 0));
+  test4(null, true, 0);
+  test4(null, false, 0);
+}