Split getClosureRepresentationInfo into MemberEntity and (ir/ast) nodes

pkg/compiler/lib/src/js_model/closure.dart

 // Copyright (c) 2017, 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.
 
 import 'package:kernel/ast.dart' as ir;
 
 import '../closure.dart';
 import '../common.dart';
 import '../common/tasks.dart';
 import '../constants/expressions.dart';
@@ skipping 52 matching lines @@
 /// http://matt.might.net/articles/closure-conversion/.
 // TODO(efortuna): Change inheritance hierarchy so that the
 // ClosureConversionTask doesn't inherit from ClosureTask because it's just a
 // glorified timer.
 class KernelClosureConversionTask extends ClosureConversionTask<ir.Node> {
   final KernelToElementMapForBuilding _elementMap;
   final GlobalLocalsMap _globalLocalsMap;
   final Map<MemberEntity, ScopeModel> _closureModels;
 
   /// Map of the scoping information that corresponds to a particular entity.
-  Map<Entity, ScopeInfo> _scopeMap = <Entity, ScopeInfo>{};
+  Map<MemberEntity, ScopeInfo> _scopeMap = <MemberEntity, ScopeInfo>{};
   Map<ir.Node, CapturedScope> _capturedScopesMap = <ir.Node, CapturedScope>{};
 
-  Map<Entity, ClosureRepresentationInfo> _closureRepresentationMap =
-      <Entity, ClosureRepresentationInfo>{};
+  Map<MemberEntity, ClosureRepresentationInfo> _memberClosureRepresentationMap =
+      <MemberEntity, ClosureRepresentationInfo>{};
+
+  // The key is either a [ir.FunctionDeclaration] or [ir.FunctionExpression].
+  Map<ir.Node, ClosureRepresentationInfo> _localClosureRepresentationMap =
+      <ir.Node, ClosureRepresentationInfo>{};
 
   KernelClosureConversionTask(Measurer measurer, this._elementMap,
       this._globalLocalsMap, this._closureModels)
       : super(measurer);
 
   /// The combined steps of generating our intermediate representation of
   /// closures that need to be rewritten and generating the element model.
   /// Ultimately these two steps will be split apart with the second step
   /// happening later in compilation just before codegen. These steps are
   /// combined here currently to provide a consistent interface to the rest of
@@ skipping 43 matching lines @@
       MemberEntity member,
       ir.FunctionNode node,
       KernelScopeInfo info,
       JsClosedWorld closedWorldRefiner) {
     KernelToLocalsMap localsMap = _globalLocalsMap.getLocalsMap(member);
     KernelClosureClass closureClass = closedWorldRefiner.buildClosureClass(
         member, node, member.library, info, node.location, localsMap);
 
     // We want the original declaration where that function is used to point
     // to the correct closure class.
-    _closureRepresentationMap[closureClass.callMethod] = closureClass;
-    Entity entity;
+    _memberClosureRepresentationMap[closureClass.callMethod] = closureClass;
     if (node.parent is ir.Member) {
-      entity = _elementMap.getMember(node.parent);
+      assert(_elementMap.getMember(node.parent) == member);
+      _memberClosureRepresentationMap[member] = closureClass;
     } else {
-      entity = localsMap.getLocalFunction(node.parent);
+      assert(node.parent is ir.FunctionExpression ||
+          node.parent is ir.FunctionDeclaration);
+      _localClosureRepresentationMap[node.parent] = closureClass;
     }
-    assert(entity != null);
-    _closureRepresentationMap[entity] = closureClass;
     return closureClass;
   }
 
   @override
-  ScopeInfo getScopeInfo(Entity entity) {
+  ScopeInfo getScopeInfo(MemberEntity entity) {
     // TODO(johnniwinther): Remove this check when constructor bodies a created
     // eagerly with the J-model; a constructor body should have it's own
     // [ClosureRepresentationInfo].
     if (entity is ConstructorBodyEntity) {
       ConstructorBodyEntity constructorBody = entity;
       entity = constructorBody.constructor;
     }
 
-    return _scopeMap[entity] ?? getClosureRepresentationInfo(entity);
+    return _scopeMap[entity] ?? getClosureInfoForMember(entity);
   }
 
   // TODO(efortuna): Eventually capturedScopesMap[node] should always
   // be non-null, and we should just test that with an assert.
   @override
   CapturedScope getCapturedScope(MemberEntity entity) {
     MemberDefinition definition = _elementMap.getMemberDefinition(entity);
     switch (definition.kind) {
       case MemberKind.regular:
       case MemberKind.constructor:
       case MemberKind.constructorBody:
       case MemberKind.closureCall:
         return _capturedScopesMap[definition.node] ?? const CapturedScope();
       default:
         throw failedAt(entity, "Unexpected member definition $definition");
     }
   }
 
   @override
   // TODO(efortuna): Eventually capturedScopesMap[node] should always
   // be non-null, and we should just test that with an assert.
   CapturedLoopScope getCapturedLoopScope(ir.Node loopNode) =>
       _capturedScopesMap[loopNode] ?? const CapturedLoopScope();
 
   @override
-  ClosureRepresentationInfo getClosureRepresentationInfo(Entity entity) {
-    var closure = _closureRepresentationMap[entity];
+  ClosureRepresentationInfo getClosureInfoForMember(MemberEntity entity) {
+    var closure = _memberClosureRepresentationMap[entity];
     assert(
         closure != null,
         "Corresponding closure class not found for $entity. "
-        "Closures found for ${_closureRepresentationMap.keys}");
+        "Closures found for ${_memberClosureRepresentationMap.keys}");
     return closure;
   }
 
   @override
-  ClosureRepresentationInfo getClosureRepresentationInfoForTesting(
-      Entity member) {
-    return _closureRepresentationMap[member];
+  ClosureRepresentationInfo getClosureInfo(ir.Node node) {
+    var closure = _localClosureRepresentationMap[node];
+    assert(
+        closure != null,
+        "Corresponding closure class not found for $node. "
+        "Closures found for ${_localClosureRepresentationMap.keys}");
+    return closure;
+  }
+
+  @override
+  ClosureRepresentationInfo getClosureInfoForMemberTesting(
+      MemberEntity entity) {
+    return _memberClosureRepresentationMap[entity];
+  }
+
+  @override
+  ClosureRepresentationInfo getClosureInfoForTesting(ir.Node node) {
+    return _localClosureRepresentationMap[node];
   }
 }
 
 class KernelScopeInfo {
   final Set<ir.VariableDeclaration> localsUsedInTryOrSync;
   final bool hasThisLocal;
   final Set<ir.VariableDeclaration> boxedVariables;
   // If boxedVariables is empty, this will be null, because no variables will
   // need to be boxed.
   final NodeBox capturedVariablesAccessor;
@@ skipping 136 matching lines @@
 
   final Map<Local, JField> localToFieldMap = new Map<Local, JField>();
 
   KernelClosureClass.fromScopeInfo(
       this.closureClassEntity,
       ir.FunctionNode closureSourceNode,
       KernelScopeInfo info,
       KernelToLocalsMap localsMap)
       : closureEntity = closureSourceNode.parent is ir.Member
             ? null
+            // TODO(johnniwinther,efortuna): This is the only place we call
+            // [getLocalFunction]. Therefore the [closureEntity] doesn't need
+            // to be derived from the node.
+            //
+            // What we should do instead: If `closureSourceNode.parent` is
+            // an [ir.FunctionDeclaration] we should use the local for its
+            // variable. If `closureSourceNode.parent` is an
+            // [ir.FunctionExpression], we should create a fresh local.
             : localsMap.getLocalFunction(closureSourceNode.parent),
         thisLocal =
             info.hasThisLocal ? new ThisLocal(localsMap.currentMember) : null,
         super.from(info, localsMap);
 
   List<Local> get createdFieldEntities => localToFieldMap.keys.toList();
 
   FieldEntity get thisFieldEntity => localToFieldMap[thisLocal];
 
   void forEachCapturedVariable(f(Local from, JField to)) {
@@ skipping 175 matching lines @@
   KernelScopeInfo scopeInfo;
 
   /// Collected [CapturedScope] data for nodes.
   Map<ir.Node, KernelCapturedScope> capturedScopesMap =
       <ir.Node, KernelCapturedScope>{};
 
   /// Collected [ScopeInfo] data for nodes.
   Map<ir.FunctionNode, KernelScopeInfo> closuresToGenerate =
       <ir.FunctionNode, KernelScopeInfo>{};
 }