Move source out of PotentialModification.

Rather than have each PotentialModification keep track of its own
source, it's less wasteful to keep track of a map from each source to
a list of the potential modifications for that source.

Change-Id: Ib579f5159100195496a850fbf795b073fc927110
Reviewed-on: https://dart-review.googlesource.com/c/93466
Reviewed-by: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Konstantin Shcheglov <scheglov@google.com>
Commit-Queue: Paul Berry <paulberry@google.com>
diff --git a/pkg/analysis_server/lib/src/nullability/constraint_gatherer.dart b/pkg/analysis_server/lib/src/nullability/constraint_gatherer.dart
index 45fc46a..d0802ad 100644
--- a/pkg/analysis_server/lib/src/nullability/constraint_gatherer.dart
+++ b/pkg/analysis_server/lib/src/nullability/constraint_gatherer.dart
@@ -427,7 +427,7 @@
         if (expression != null) {
           checkNotNull = CheckExpression(expression);
           _variables.recordExpressionChecks(
-              expression, ExpressionChecks(_source, checkNotNull));
+              _source, expression, ExpressionChecks(checkNotNull));
         }
         // nullable_src => nullable_dst | check_expr
         _recordFact(ConstraintVariable.or(
diff --git a/pkg/analysis_server/lib/src/nullability/constraint_variable_gatherer.dart b/pkg/analysis_server/lib/src/nullability/constraint_variable_gatherer.dart
index f722bad..f917e72 100644
--- a/pkg/analysis_server/lib/src/nullability/constraint_variable_gatherer.dart
+++ b/pkg/analysis_server/lib/src/nullability/constraint_variable_gatherer.dart
@@ -149,10 +149,9 @@
     var nullable = node.question == null
         ? TypeIsNullable(node.end)
         : ConstraintVariable.always;
-    var decoratedType = DecoratedTypeAnnotation(
-        type, nullable, _source, node.end,
+    var decoratedType = DecoratedTypeAnnotation(type, nullable, node.end,
         typeArguments: typeArguments);
-    _variables.recordDecoratedTypeAnnotation(node, decoratedType);
+    _variables.recordDecoratedTypeAnnotation(_source, node, decoratedType);
     return decoratedType;
   }
 
@@ -193,7 +192,7 @@
 
   /// Associates decorated type information with the given [type] node.
   void recordDecoratedTypeAnnotation(
-      TypeAnnotation node, DecoratedTypeAnnotation type);
+      Source source, TypeAnnotation node, DecoratedTypeAnnotation type);
 
   /// Associates a constraint variable with the question of whether the given
   /// named parameter should be optional (should not have a `required`
@@ -232,5 +231,6 @@
   void recordDecoratedExpressionType(Expression node, DecoratedType type);
 
   /// Associates a set of nullability checks with the given expression [node].
-  void recordExpressionChecks(Expression expression, ExpressionChecks checks);
+  void recordExpressionChecks(
+      Source source, Expression expression, ExpressionChecks checks);
 }
diff --git a/pkg/analysis_server/lib/src/nullability/decorated_type.dart b/pkg/analysis_server/lib/src/nullability/decorated_type.dart
index 53c023b..2dde11f 100644
--- a/pkg/analysis_server/lib/src/nullability/decorated_type.dart
+++ b/pkg/analysis_server/lib/src/nullability/decorated_type.dart
@@ -7,7 +7,6 @@
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
 import 'package:analyzer/src/dart/element/type.dart';
-import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer_plugin/protocol/protocol_common.dart' show SourceEdit;
 
 /// Representation of a type in the code to be migrated.  In addition to
@@ -178,13 +177,10 @@
 /// the source code to reflect its nullability.
 class DecoratedTypeAnnotation extends DecoratedType
     implements PotentialModification {
-  @override
-  final Source source;
-
   final int _offset;
 
   DecoratedTypeAnnotation(
-      DartType type, ConstraintVariable nullable, this.source, this._offset,
+      DartType type, ConstraintVariable nullable, this._offset,
       {List<DecoratedType> typeArguments = const []})
       : super(type, nullable, typeArguments: typeArguments);
 
diff --git a/pkg/analysis_server/lib/src/nullability/expression_checks.dart b/pkg/analysis_server/lib/src/nullability/expression_checks.dart
index 3c4af10..14f47db 100644
--- a/pkg/analysis_server/lib/src/nullability/expression_checks.dart
+++ b/pkg/analysis_server/lib/src/nullability/expression_checks.dart
@@ -3,7 +3,6 @@
 // BSD-style license that can be found in the LICENSE file.
 
 import 'package:analysis_server/src/nullability/transitional_api.dart';
-import 'package:analyzer/src/generated/source.dart';
 import 'package:analyzer_plugin/protocol/protocol_common.dart';
 
 /// Container for information gathered during nullability migration about the
@@ -14,14 +13,11 @@
 /// that the expression is not null.  We need to add other checks, e.g. to check
 /// that a List<int?> is actually a List<int>.
 class ExpressionChecks extends PotentialModification {
-  @override
-  final Source source;
-
   /// Constraint variable whose value will be `true` if this expression requires
   /// a null check.
   final CheckExpression nullCheck;
 
-  ExpressionChecks(this.source, this.nullCheck);
+  ExpressionChecks(this.nullCheck);
 
   @override
   bool get isEmpty => !nullCheck.value;
diff --git a/pkg/analysis_server/lib/src/nullability/provisional_api.dart b/pkg/analysis_server/lib/src/nullability/provisional_api.dart
index 6352b4d..dd8ae46 100644
--- a/pkg/analysis_server/lib/src/nullability/provisional_api.dart
+++ b/pkg/analysis_server/lib/src/nullability/provisional_api.dart
@@ -77,9 +77,12 @@
             permissive: permissive, assumptions: assumptions);
 
   void finish() {
-    _analyzerMigration.finish().forEach((pm) {
-      listener.addFix(_SingleNullabilityFix(pm));
-    });
+    for (var entry in _analyzerMigration.finish().entries) {
+      var source = entry.key;
+      for (var potentialModification in entry.value) {
+        listener.addFix(_SingleNullabilityFix(source, potentialModification));
+      }
+    }
   }
 
   void prepareInput(ResolvedUnitResult result) {
@@ -129,7 +132,7 @@
   final NullabilityFixKind kind;
 
   factory _SingleNullabilityFix(
-      analyzer.PotentialModification potentialModification) {
+      Source source, analyzer.PotentialModification potentialModification) {
     // TODO(paulberry): once everything is migrated into the analysis server,
     // the migration engine can just create SingleNullabilityFix objects
     // directly and set their kind appropriately; we won't need to translate the
@@ -148,8 +151,8 @@
     } else {
       throw new UnimplementedError('TODO(paulberry)');
     }
-    return _SingleNullabilityFix._(potentialModification.modifications.toList(),
-        potentialModification.source, kind);
+    return _SingleNullabilityFix._(
+        potentialModification.modifications.toList(), source, kind);
   }
 
   _SingleNullabilityFix._(this.sourceEdits, this.source, this.kind);
diff --git a/pkg/analysis_server/lib/src/nullability/transitional_api.dart b/pkg/analysis_server/lib/src/nullability/transitional_api.dart
index ba905db..ff2a53f 100644
--- a/pkg/analysis_server/lib/src/nullability/transitional_api.dart
+++ b/pkg/analysis_server/lib/src/nullability/transitional_api.dart
@@ -41,11 +41,7 @@
 
   final _KeepNode elseStatement;
 
-  @override
-  final Source source;
-
-  factory ConditionalModification(
-      Source source, AstNode node, ConditionalDiscard discard) {
+  factory ConditionalModification(AstNode node, ConditionalDiscard discard) {
     if (node is IfStatement) {
       return ConditionalModification._(
           node.offset,
@@ -54,22 +50,14 @@
           discard,
           _KeepNode(node.condition),
           _KeepNode(node.thenStatement),
-          _KeepNode(node.elseStatement),
-          source);
+          _KeepNode(node.elseStatement));
     } else {
       throw new UnimplementedError('TODO(paulberry)');
     }
   }
 
-  ConditionalModification._(
-      this.offset,
-      this.end,
-      this.isStatement,
-      this.discard,
-      this.condition,
-      this.thenStatement,
-      this.elseStatement,
-      this.source);
+  ConditionalModification._(this.offset, this.end, this.isStatement,
+      this.discard, this.condition, this.thenStatement, this.elseStatement);
 
   @override
   bool get isEmpty => discard.keepTrue.value && discard.keepFalse.value;
@@ -171,7 +159,7 @@
       this.assumptions: const NullabilityMigrationAssumptions()})
       : _permissive = permissive;
 
-  List<PotentialModification> finish() {
+  Map<Source, List<PotentialModification>> finish() {
     _constraints.applyHeuristics();
     return _variables.getPotentialModifications();
   }
@@ -211,15 +199,12 @@
 /// Records information about the possible addition of a `@required` annotation
 /// to the source code.
 class PotentiallyAddRequired extends PotentialModification {
-  @override
-  final Source source;
-
   final ConstraintVariable _optionalVariable;
 
   final int _offset;
 
   PotentiallyAddRequired(
-      this.source, DefaultFormalParameter parameter, this._optionalVariable)
+      DefaultFormalParameter parameter, this._optionalVariable)
       : _offset = parameter.offset;
 
   @override
@@ -238,14 +223,12 @@
   /// Gets the individual migrations that need to be done, considering the
   /// solution to the constraint equations.
   Iterable<SourceEdit> get modifications;
-
-  Source get source;
 }
 
 class Variables implements VariableRecorder, VariableRepository {
   final _decoratedElementTypes = <Element, DecoratedType>{};
 
-  final _potentialModifications = <PotentialModification>[];
+  final _potentialModifications = <Source, List<PotentialModification>>{};
 
   @override
   DecoratedType decoratedElementType(Element element, {bool create: false}) =>
@@ -253,14 +236,14 @@
           ? DecoratedType.forElement(element)
           : throw StateError('No element found');
 
-  List<PotentialModification> getPotentialModifications() =>
-      _potentialModifications.where((m) => !m.isEmpty).toList();
+  Map<Source, List<PotentialModification>> getPotentialModifications() =>
+      _potentialModifications;
 
   @override
   void recordConditionalDiscard(
       Source source, AstNode node, ConditionalDiscard conditionalDiscard) {
-    _potentialModifications
-        .add(ConditionalModification(source, node, conditionalDiscard));
+    _addPotentialModification(
+        source, ConditionalModification(node, conditionalDiscard));
   }
 
   void recordDecoratedElementType(Element element, DecoratedType type) {
@@ -270,20 +253,26 @@
   void recordDecoratedExpressionType(Expression node, DecoratedType type) {}
 
   void recordDecoratedTypeAnnotation(
-      TypeAnnotation node, DecoratedTypeAnnotation type) {
-    _potentialModifications.add(type);
+      Source source, TypeAnnotation node, DecoratedTypeAnnotation type) {
+    _addPotentialModification(source, type);
   }
 
   @override
-  void recordExpressionChecks(Expression expression, ExpressionChecks checks) {
-    _potentialModifications.add(checks);
+  void recordExpressionChecks(
+      Source source, Expression expression, ExpressionChecks checks) {
+    _addPotentialModification(source, checks);
   }
 
   @override
   void recordPossiblyOptional(Source source, DefaultFormalParameter parameter,
       ConstraintVariable variable) {
-    _potentialModifications
-        .add(PotentiallyAddRequired(source, parameter, variable));
+    _addPotentialModification(
+        source, PotentiallyAddRequired(parameter, variable));
+  }
+
+  void _addPotentialModification(
+      Source source, PotentialModification potentialModification) {
+    (_potentialModifications[source] ??= []).add(potentialModification);
   }
 }
 
diff --git a/pkg/analysis_server/test/src/nullability/migration_visitor_test.dart b/pkg/analysis_server/test/src/nullability/migration_visitor_test.dart
index bd9df3b..e9db0d5 100644
--- a/pkg/analysis_server/test/src/nullability/migration_visitor_test.dart
+++ b/pkg/analysis_server/test/src/nullability/migration_visitor_test.dart
@@ -1024,14 +1024,16 @@
     _decoratedExpressionTypes[_normalizeExpression(node)] = type;
   }
 
-  void recordDecoratedTypeAnnotation(TypeAnnotation node, DecoratedType type) {
-    super.recordDecoratedTypeAnnotation(node, type);
+  void recordDecoratedTypeAnnotation(
+      Source source, TypeAnnotation node, DecoratedType type) {
+    super.recordDecoratedTypeAnnotation(source, node, type);
     _decoratedTypeAnnotations[node] = type;
   }
 
   @override
-  void recordExpressionChecks(Expression expression, ExpressionChecks checks) {
-    super.recordExpressionChecks(expression, checks);
+  void recordExpressionChecks(
+      Source source, Expression expression, ExpressionChecks checks) {
+    super.recordExpressionChecks(source, expression, checks);
     _expressionChecks[_normalizeExpression(expression)] = checks;
   }