[vm/aot/tfa] Whole-program propagation of closure values

This change adds propagation of closure values to TFA.
For now, inferred closure values are used only when they are
constant (tear-off of a static method or a constructor).

Arbitrary closures can now be referenced from unrelated
members via closure-id metadata.

In addition, this change fixes an incorrect stack trace when
an implicit closure (tear-off) was propagated and its call was
inlined. Inlining interval was not recorded because of the missing
source position of a call to a target method within implicit closure.

TEST=pkg/vm/testcases/transformations/type_flow/transformer/closures.dart
Issue: #39692

Change-Id: I3590da91b6057e0b55a8614382dba1bbcc267b39
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/325447
Commit-Queue: Alexander Markov <[email protected]>
Reviewed-by: Martin Kustermann <[email protected]>
Reviewed-by: Slava Egorov <[email protected]>

Author: alexmarkov

pkg/dart2wasm/lib/target.dart

@@ -102,6 +102,7 @@ class WasmTarget extends Target {
   Class? _wasmImmutableSet;
   Class? _oneByteString;
   Class? _twoByteString;
+  Class? _closure;
   Map<String, Class>? _nativeClasses;
 
   @override
@@ -479,6 +480,11 @@ class WasmTarget extends Target {
         coreTypes.index.getClass('dart:_string', 'OneByteString');
   }
 
+  @override
+  Class concreteClosureClass(CoreTypes coreTypes) {
+    return _closure ??= coreTypes.index.getClass('dart:core', '_Closure');
+  }
+
   @override
   bool isSupportedPragma(String pragmaName) => pragmaName.startsWith("wasm:");
 

pkg/kernel/lib/target/targets.dart

@@ -515,6 +515,7 @@ abstract class Target {
   Class? concreteConstMapLiteralClass(CoreTypes coreTypes) => null;
   Class? concreteSetLiteralClass(CoreTypes coreTypes) => null;
   Class? concreteConstSetLiteralClass(CoreTypes coreTypes) => null;
+  Class? concreteClosureClass(CoreTypes coreTypes) => null;
   Class getRecordImplementationClass(CoreTypes coreTypes,
           int numPositionalFields, List<String> namedFields) =>
       throw UnsupportedError('Target.getRecordImplementationClass');

pkg/vm/bin/dump_kernel.dart

@@ -8,6 +8,7 @@ import 'package:kernel/kernel.dart' show Component, writeComponentToText;
 import 'package:kernel/binary/ast_from_binary.dart'
     show BinaryBuilderWithMetadata;
 
+import 'package:vm/metadata/closure_id.dart' show ClosureIdMetadataRepository;
 import 'package:vm/metadata/direct_call.dart' show DirectCallMetadataRepository;
 import 'package:vm/metadata/inferred_type.dart'
     show InferredTypeMetadataRepository, InferredArgTypeMetadataRepository;
@@ -50,6 +51,7 @@ main(List<String> arguments) async {
   component.addMetadataRepository(new UnreachableNodeMetadataRepository());
   component.addMetadataRepository(new CallSiteAttributesMetadataRepository());
   component.addMetadataRepository(new LoadingUnitsMetadataRepository());
+  component.addMetadataRepository(new ClosureIdMetadataRepository());
 
   final List<int> bytes = new File(input).readAsBytesSync();
   new BinaryBuilderWithMetadata(bytes).readComponent(component);

pkg/vm/lib/metadata/closure_id.dart

@@ -0,0 +1,71 @@
+// 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.
+
+import 'package:kernel/ast.dart';
+
+/// Repository for persistent closure IDs.
+class ClosureIdMetadataRepository extends MetadataRepository<int> {
+  static const String repositoryTag = 'vm.closure-id';
+
+  @override
+  String get tag => repositoryTag;
+
+  // For a LocalFunction: id within an enclosing Member,
+  // with 0 reserved for the tear-off of the Member.
+  //
+  // For a Member: number of nested closures.
+  @override
+  final Map<TreeNode, int> mapping = {};
+
+  @override
+  void writeToBinary(int metadata, Node node, BinarySink sink) {
+    sink.writeUInt30(metadata);
+  }
+
+  @override
+  int readFromBinary(Node node, BinarySource source) {
+    return source.readUInt30();
+  }
+
+  /// Return closure ID within the enclosing [Member].
+  ///
+  /// Closures should be indexed within enclosing [Member]
+  /// using [indexClosures].
+  int getClosureId(LocalFunction closure) => mapping[closure]!;
+
+  /// Assign IDs for all closures within [member].
+  void indexClosures(Member member) {
+    if (mapping.containsKey(member)) {
+      return;
+    }
+    _ClosureIndexer indexer = _ClosureIndexer(this, member);
+    member.accept(indexer);
+    mapping[member] = indexer.index - _ClosureIndexer.firstClosureIndex;
+  }
+}
+
+class _ClosureIndexer extends RecursiveVisitor<void> {
+  // Zero is reserved for tear-offs.
+  static int firstClosureIndex = 1;
+
+  final ClosureIdMetadataRepository _repository;
+  final Member member;
+  int index = firstClosureIndex;
+
+  _ClosureIndexer(this._repository, this.member);
+
+  @override
+  void visitFunctionDeclaration(FunctionDeclaration node) =>
+      _visitLocalFunction(node);
+
+  @override
+  void visitFunctionExpression(FunctionExpression node) =>
+      _visitLocalFunction(node);
+
+  void _visitLocalFunction(LocalFunction node) {
+    assert(index > 0);
+    _repository.mapping[node] = index++;
+    node.visitChildren(this);
+  }
+}

pkg/vm/lib/metadata/inferred_type.dart

@@ -11,6 +11,8 @@ import 'package:kernel/src/printer.dart';
 class InferredType {
   final Reference? _concreteClassReference;
   final Constant? _constantValue;
+  final Reference? _closureMemberReference;
+  final int _closureId;
   final int _flags;
 
   static const int flagNullable = 1 << 0;
@@ -19,10 +21,14 @@ class InferredType {
   // For invocations: whether to use the unchecked entry-point.
   static const int flagSkipCheck = 1 << 2;
 
+  // Contains inferred constant value.
   static const int flagConstant = 1 << 3;
 
   static const int flagReceiverNotInt = 1 << 4;
 
+  // Contains inferred closure value.
+  static const int flagClosure = 1 << 5;
+
   // Entire list may be null if no type arguments were inferred.
   // Will always be null if `concreteClass` is null.
   //
@@ -33,31 +39,44 @@ class InferredType {
   // argument (in the runtime type) is always exactly a particular `DartType`.
   final List<DartType?>? exactTypeArguments;
 
-  InferredType(
-      Class? concreteClass, bool nullable, bool isInt, Constant? constantValue,
+  InferredType(Class? concreteClass, bool nullable, bool isInt,
+      Constant? constantValue, Member? closureMember, int closureId,
       {List<DartType?>? exactTypeArguments,
       bool skipCheck = false,
       bool receiverNotInt = false})
       : this._byReference(
             concreteClass?.reference,
             constantValue,
+            closureMember?.reference,
+            closureId,
             (nullable ? flagNullable : 0) |
                 (isInt ? flagInt : 0) |
                 (skipCheck ? flagSkipCheck : 0) |
                 (constantValue != null ? flagConstant : 0) |
-                (receiverNotInt ? flagReceiverNotInt : 0),
+                (receiverNotInt ? flagReceiverNotInt : 0) |
+                (closureMember != null ? flagClosure : 0),
             exactTypeArguments);
 
-  InferredType._byReference(this._concreteClassReference, this._constantValue,
-      this._flags, this.exactTypeArguments) {
+  InferredType._byReference(
+      this._concreteClassReference,
+      this._constantValue,
+      this._closureMemberReference,
+      this._closureId,
+      this._flags,
+      this.exactTypeArguments) {
     assert(exactTypeArguments == null || _concreteClassReference != null);
     assert(_constantValue == null || _concreteClassReference != null);
+    assert(_closureMemberReference == null || _concreteClassReference != null);
+    assert(_closureId >= 0);
   }
 
   Class? get concreteClass => _concreteClassReference?.asClass;
 
   Constant? get constantValue => _constantValue;
 
+  Member? get closureMember => _closureMemberReference?.asMember;
+  int get closureId => _closureId;
+
   bool get nullable => (_flags & flagNullable) != 0;
   bool get isInt => (_flags & flagInt) != 0;
   bool get skipCheck => (_flags & flagSkipCheck) != 0;
@@ -96,6 +115,10 @@ class InferredType {
     if (receiverNotInt) {
       buf.write(' (receiver not int)');
     }
+    if (closureMember != null) {
+      buf.write(
+          ' (closure ${closureId} in ${closureMember!.toText(astTextStrategyForTesting)})');
+    }
     return buf.toString();
   }
 }
@@ -116,10 +139,16 @@ class InferredTypeMetadataRepository extends MetadataRepository<InferredType> {
     // them for optimizations.
     sink.writeNullAllowedCanonicalNameReference(
         metadata.concreteClass?.reference);
+    final flags = metadata._flags;
     sink.writeByte(metadata._flags);
-    if (metadata.constantValue != null) {
+    if ((flags & InferredType.flagConstant) != 0) {
       sink.writeConstantReference(metadata.constantValue!);
     }
+    if ((flags & InferredType.flagClosure) != 0) {
+      sink.writeNullAllowedCanonicalNameReference(
+          metadata.closureMember!.reference);
+      sink.writeUInt30(metadata.closureId);
+    }
   }
 
   @override
@@ -132,8 +161,13 @@ class InferredTypeMetadataRepository extends MetadataRepository<InferredType> {
     final constantValue = (flags & InferredType.flagConstant) != 0
         ? source.readConstantReference()
         : null;
-    return new InferredType._byReference(
-        concreteClassReference, constantValue, flags, null);
+    final closureMemberReference = (flags & InferredType.flagClosure) != 0
+        ? source.readNullableCanonicalNameReference()!.reference
+        : null;
+    final closureId =
+        (flags & InferredType.flagClosure) != 0 ? source.readUInt30() : 0;
+    return new InferredType._byReference(concreteClassReference, constantValue,
+        closureMemberReference, closureId, flags, null);
   }
 }
 

pkg/vm/lib/target/vm.dart

@@ -39,6 +39,7 @@ class VmTarget extends Target {
   Class? _twoByteString;
   Class? _smi;
   Class? _double; // _Double, not double.
+  Class? _closure;
   Class? _syncStarIterable;
 
   VmTarget(this.flags);
@@ -490,6 +491,11 @@ class VmTarget extends Target {
         coreTypes.index.getClass('dart:core', '_OneByteString');
   }
 
+  @override
+  Class concreteClosureClass(CoreTypes coreTypes) {
+    return _closure ??= coreTypes.index.getClass('dart:core', '_Closure');
+  }
+
   @override
   Class? concreteAsyncResultClass(CoreTypes coreTypes) =>
       coreTypes.futureImplClass;

pkg/vm/lib/transformations/call_site_annotator.dart

@@ -83,19 +83,6 @@ class AnnotateWithStaticTypes extends RecursiveVisitor {
     }
   }
 
-  @override
-  visitFunctionInvocation(FunctionInvocation node) {
-    super.visitFunctionInvocation(node);
-
-    final DartType receiverType =
-        node.receiver.getStaticType(_staticTypeContext!);
-    if (receiverType is FunctionType &&
-        node.kind == FunctionAccessKind.Function) {
-      throw 'Node ${node.runtimeType}: $node at ${node.location} has receiver'
-          ' static type $receiverType, but kind ${node.kind}';
-    }
-  }
-
   @override
   visitEqualsCall(EqualsCall node) {
     super.visitEqualsCall(node);

pkg/vm/lib/transformations/type_flow/analysis.dart

@@ -315,9 +315,29 @@ final class _DirectInvocation extends _Invocation {
       if (selector.callKind == CallKind.PropertyGet) {
         // Tear-off.
         // TODO(alexmarkov): capture receiver type
-        assert((member is Procedure) && !member.isGetter && !member.isSetter);
+        assert((member is Procedure) &&
+            !member.isGetter &&
+            !member.isSetter &&
+            !member.isFactory &&
+            !member.isAbstract);
         typeFlowAnalysis.addRawCall(new DirectSelector(member));
         typeFlowAnalysis._tearOffTaken.add(member);
+        final Class? concreteClass = typeFlowAnalysis.target
+            .concreteClosureClass(typeFlowAnalysis.coreTypes);
+        if (concreteClass != null) {
+          if (!member.isInstanceMember) {
+            return typeFlowAnalysis
+                .addAllocatedClass(concreteClass)
+                .cls
+                .constantConcreteType(
+                    StaticTearOffConstant(member as Procedure));
+          } else {
+            return typeFlowAnalysis
+                .addAllocatedClass(concreteClass)
+                .cls
+                .closureConcreteType(member, null);
+          }
+        }
         return nullableAnyType;
       } else {
         // Call via getter.

pkg/vm/lib/transformations/type_flow/signature_shaking.dart

@@ -182,7 +182,7 @@ class _ParameterInfo {
     // constant value in every implementation. The constant value inferred does
     // not have to be the same across implementations, as it is specialized in
     // each implementation individually.
-    if (!(type is ConcreteType && type.constant != null ||
+    if (!(type is ConcreteType && type.attributes?.constant != null ||
         type is NullableType && type.baseType is EmptyType)) {
       isConstant = false;
     }
@@ -372,7 +372,7 @@ class _Transform extends RecursiveVisitor {
     if (param.isConstant) {
       Type type = shaker.typeFlowAnalysis.argumentType(member, variable)!;
       if (type is ConcreteType) {
-        value = type.constant!;
+        value = type.attributes!.constant!;
       } else {
         assert(type is NullableType && type.baseType is EmptyType);
         value = NullConstant();
@@ -499,7 +499,8 @@ class _Transform extends RecursiveVisitor {
   void visitVariableGet(VariableGet node) {
     Constant? constantValue = eliminatedParams[node.variable];
     if (constantValue != null) {
-      node.replaceWith(ConstantExpression(constantValue));
+      node.replaceWith(ConstantExpression(
+          constantValue, constantValue.getType(typeContext)));
     }
   }