Record instances which are not annotations

There is no need to restrict the recording to annotations.
It also opens up more interesting use cases, for example recording instances of `const StringAsset('id').load()` if the class `StringAsset` has the `RecordUse` annotation.

Also includes a small refactoring in a separate commit.

TESTED=pkg/vm/testcases/transformations/record_use/instance_not_annotation.dart
Change-Id: I262995760a8ba8bc14c9cf01e5eb1b47f1278282
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/387700
Reviewed-by: Martin Kustermann <[email protected]>
Commit-Queue: Moritz Sümmermann <[email protected]>

Author: mosuem

pkg/vm/lib/transformations/record_use/constant_collector.dart

@@ -0,0 +1,140 @@
+// 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.
+
+import 'package:kernel/ast.dart';
+import 'package:front_end/src/kernel/record_use.dart' as recordUse;
+
+/// Expose only the [collect] method of a [_ConstantCollector] to outside use.
+extension type ConstantCollector(_ConstantCollector _collector) {
+  ConstantCollector.collectWith(
+      Function(
+        ConstantExpression context,
+        InstanceConstant constant,
+      ) collector)
+      : _collector = _ConstantCollector(collector);
+
+  void collect(ConstantExpression expression) => _collector.collect(expression);
+}
+
+/// A visitor traversing constants and storing instance constants with the
+/// `@RecordUse` annotation using the [collector] callback.
+class _ConstantCollector implements ConstantVisitor {
+  /// The collector callback which records the constant.
+  final void Function(
+    ConstantExpression context,
+    InstanceConstant constant,
+  ) collector;
+
+  /// The expression in which the constant was found.
+  ConstantExpression? _expression;
+
+  /// A cache to avoid having to re-check for annotations.
+  final Map<Class, bool> _hasRecordUseAnnotation = {};
+
+  final Set<Constant> _visited = {};
+
+  _ConstantCollector(this.collector);
+
+  void collect(ConstantExpression node) {
+    _expression = node;
+    handleConstantReference(node.constant);
+    _expression = null;
+  }
+
+  void handleConstantReference(Constant constant) {
+    if (_visited.add(constant)) {
+      constant.accept(this);
+    }
+  }
+
+  @override
+  void visitListConstant(ListConstant constant) {
+    for (final entry in constant.entries) {
+      handleConstantReference(entry);
+    }
+  }
+
+  @override
+  void visitMapConstant(MapConstant constant) {
+    for (final entry in constant.entries) {
+      handleConstantReference(entry.key);
+      handleConstantReference(entry.value);
+    }
+  }
+
+  @override
+  void visitSetConstant(SetConstant constant) {
+    for (final entry in constant.entries) {
+      handleConstantReference(entry);
+    }
+  }
+
+  @override
+  void visitRecordConstant(RecordConstant constant) {
+    for (final value in constant.positional) {
+      handleConstantReference(value);
+    }
+    for (final value in constant.named.values) {
+      handleConstantReference(value);
+    }
+  }
+
+  @override
+  void visitInstanceConstant(InstanceConstant constant) {
+    assert(_expression != null);
+    final classNode = constant.classNode;
+    if (_hasRecordUseAnnotation[classNode] ??=
+        recordUse.findRecordUseAnnotation(classNode).isNotEmpty) {
+      collector(_expression!, constant);
+    }
+    for (final value in constant.fieldValues.values) {
+      handleConstantReference(value);
+    }
+  }
+
+  @override
+  visitAuxiliaryConstant(AuxiliaryConstant node) {
+    throw UnsupportedError('Cannot record an `AuxiliaryConstant`.');
+  }
+
+  @override
+  visitBoolConstant(BoolConstant node) {}
+
+  @override
+  visitConstructorTearOffConstant(ConstructorTearOffConstant node) {}
+
+  @override
+  visitDoubleConstant(DoubleConstant node) {}
+
+  @override
+  visitInstantiationConstant(InstantiationConstant node) {}
+
+  @override
+  visitIntConstant(IntConstant node) {}
+
+  @override
+  visitNullConstant(NullConstant node) {}
+
+  @override
+  visitRedirectingFactoryTearOffConstant(
+      RedirectingFactoryTearOffConstant node) {}
+
+  @override
+  visitStaticTearOffConstant(StaticTearOffConstant node) {}
+
+  @override
+  visitStringConstant(StringConstant node) {}
+
+  @override
+  visitSymbolConstant(SymbolConstant node) {}
+
+  @override
+  visitTypeLiteralConstant(TypeLiteralConstant node) {}
+
+  @override
+  visitTypedefTearOffConstant(TypedefTearOffConstant node) {}
+
+  @override
+  visitUnevaluatedConstant(UnevaluatedConstant node) {}
+}

pkg/vm/lib/transformations/record_use/record_call.dart

@@ -9,14 +9,29 @@ import 'package:record_use/record_use_internal.dart';
 import 'package:vm/metadata/loading_units.dart';
 import 'package:vm/transformations/record_use/record_use.dart';
 
-class StaticCallRecorder {
-  final Map<ast.Procedure, Usage<CallReference>> callsForMethod = {};
+/// Record calls and their constant arguments. Currently tracks
+/// * static or top-level method calls through [recordStaticInvocation]
+/// * tear-offs through [recordConstantExpression]
+///
+/// The result of adding calls can be fetched from [foundCalls].
+class CallRecorder {
+  /// The collection of recorded calls found so far.
+  Iterable<Usage<CallReference>> get foundCalls => _callsForMethod.values;
+
+  /// Keep track of the calls which are recorded, to easily add newly found
+  /// ones.
+  final Map<ast.Procedure, Usage<CallReference>> _callsForMethod = {};
+
+  /// The ordered list of loading units to retrieve the loading unit index from.
   final List<LoadingUnit> _loadingUnits;
-  final Uri source;
 
-  StaticCallRecorder(this.source, this._loadingUnits);
+  /// The source uri to base relative URIs off of.
+  final Uri _source;
 
-  void recordStaticCall(ast.StaticInvocation node) {
+  CallRecorder(this._source, this._loadingUnits);
+
+  /// Will record a static invocation if it is annotated with `@RecordUse`.
+  void recordStaticInvocation(ast.StaticInvocation node) {
     final annotations = recordUse.findRecordUseAnnotation(node.target);
     if (annotations.isNotEmpty) {
       final call = _getCall(node.target);
@@ -26,18 +41,27 @@ class StaticCallRecorder {
     }
   }
 
-  void recordTearoff(ast.ConstantExpression node) {
+  /// Will record a tear-off if the target is annotated with `@RecordUse`.
+  void recordConstantExpression(ast.ConstantExpression node) {
     final constant = node.constant;
     if (constant is ast.StaticTearOffConstant) {
-      final annotations = recordUse.findRecordUseAnnotation(constant.target);
-      if (annotations.isNotEmpty) {
-        final call = _getCall(constant.target);
-        CallReference reference = _collectTearOff(constant, node);
-        call.references.add(reference);
+      final hasRecordUseAnnotation =
+          recordUse.findRecordUseAnnotation(constant.target).isNotEmpty;
+      if (hasRecordUseAnnotation) {
+        _recordTearOff(constant, node);
       }
     }
   }
 
+  void _recordTearOff(
+    ast.StaticTearOffConstant constant,
+    ast.ConstantExpression node,
+  ) {
+    final call = _getCall(constant.target);
+    final reference = _collectTearOff(constant, node);
+    call.references.add(reference);
+  }
+
   /// Record a tear off as a call with all non-const arguments.
   CallReference _collectTearOff(
     ast.StaticTearOffConstant constant,
@@ -53,7 +77,7 @@ class StaticCallRecorder {
       ),
     );
     return CallReference(
-      location: node.location!.recordLocation(source),
+      location: node.location!.recordLocation(_source),
       arguments: Arguments(nonConstArguments: nonConstArguments),
     );
   }
@@ -62,10 +86,8 @@ class StaticCallRecorder {
   /// shared across multiple calls to the same method.
   Usage _getCall(ast.Procedure target) {
     final definition = _definitionFromMember(target);
-    return callsForMethod.putIfAbsent(
-      target,
-      () => Usage(definition: definition, references: []),
-    );
+    return _callsForMethod[target] ??=
+        Usage(definition: definition, references: []);
   }
 
   CallReference _createCallReference(ast.StaticInvocation node) {
@@ -89,7 +111,7 @@ class StaticCallRecorder {
     final namedGrouped = _groupByNull(namedArguments);
 
     return CallReference(
-      location: node.location!.recordLocation(source),
+      location: node.location!.recordLocation(_source),
       loadingUnit: loadingUnitForNode(node, _loadingUnits).toString(),
       arguments: Arguments(
         constArguments: ConstArguments(
@@ -126,15 +148,15 @@ class StaticCallRecorder {
 
   Definition _definitionFromMember(ast.Member target) {
     final enclosingLibrary = target.enclosingLibrary;
-    String file = getImportUri(enclosingLibrary, source);
+    String file = getImportUri(enclosingLibrary, _source);
 
     return Definition(
       identifier: Identifier(
         importUri: file,
         parent: target.enclosingClass?.name,
         name: target.name.text,
       ),
-      location: target.location!.recordLocation(source),
+      location: target.location!.recordLocation(_source),
       loadingUnit:
           loadingUnitForNode(enclosingLibrary, _loadingUnits).toString(),
     );

pkg/vm/lib/transformations/record_use/record_instance.dart

@@ -2,78 +2,90 @@
 // 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:front_end/src/kernel/record_use.dart' as recordUse;
 import 'package:kernel/ast.dart' as ast;
 import 'package:record_use/record_use_internal.dart';
 import 'package:vm/metadata/loading_units.dart';
 import 'package:vm/transformations/record_use/record_use.dart';
 
-class InstanceUseRecorder {
-  final Map<ast.Class, Usage<InstanceReference>> instancesForClass = {};
+import 'constant_collector.dart';
+
+/// Record a const instance by calling [recordConstantExpression]. After all the
+/// const instances have been recorded, retrieve them using [foundInstances].
+class InstanceRecorder {
+  /// The collection of recorded instances found so far.
+  Iterable<Usage<InstanceReference>> get foundInstances =>
+      _instancesForClass.values;
+
+  /// Keep track of the classes which are recorded, to easily add found
+  /// instances.
+  final Map<ast.Class, Usage<InstanceReference>> _instancesForClass = {};
+
+  /// The ordered list of loading units to retrieve the loading unit index from.
   final List<LoadingUnit> _loadingUnits;
-  final Uri source;
 
-  InstanceUseRecorder(this.source, this._loadingUnits);
+  /// The source uri to base relative URIs off of.
+  final Uri _source;
+
+  /// A visitor traversing and collecting constants.
+  late final ConstantCollector collector;
 
-  void recordAnnotationUse(ast.ConstantExpression node) {
-    final constant = node.constant;
-    if (constant is ast.InstanceConstant) {
-      if (recordUse.findRecordUseAnnotation(constant.classNode).isNotEmpty) {
-        _collectUseInformation(node, constant);
-      }
-    }
+  InstanceRecorder(this._source, this._loadingUnits) {
+    collector = ConstantCollector.collectWith(_collectInstance);
   }
 
-  void _collectUseInformation(
-    ast.ConstantExpression node,
+  void recordConstantExpression(ast.ConstantExpression node) =>
+      collector.collect(node);
+
+  void _collectInstance(
+    ast.ConstantExpression expression,
     ast.InstanceConstant constant,
   ) {
     // Collect the name and definition location of the invocation. This is
     // shared across multiple calls to the same method.
     final existingInstance = _getCall(constant.classNode);
 
     // Collect the (int, bool, double, or String) arguments passed in the call.
-    existingInstance.references.add(_createInstanceReference(node, constant));
+    existingInstance.references
+        .add(_createInstanceReference(expression, constant));
   }
 
   /// Collect the name and definition location of the invocation. This is
   /// shared across multiple calls to the same method.
   Usage<InstanceReference> _getCall(ast.Class cls) {
     final definition = _definitionFromClass(cls);
-    return instancesForClass.putIfAbsent(
-      cls,
-      () => Usage(definition: definition, references: []),
-    );
+    return _instancesForClass[cls] ??=
+        Usage(definition: definition, references: []);
   }
 
   Definition _definitionFromClass(ast.Class cls) {
     final enclosingLibrary = cls.enclosingLibrary;
-    String file = getImportUri(enclosingLibrary, source);
+    final file = getImportUri(enclosingLibrary, _source);
 
     return Definition(
       identifier: Identifier(importUri: file, name: cls.name),
-      location: cls.location!.recordLocation(source),
+      location: cls.location!.recordLocation(_source),
       loadingUnit:
           loadingUnitForNode(cls.enclosingLibrary, _loadingUnits).toString(),
     );
   }
 
   InstanceReference _createInstanceReference(
-    ast.ConstantExpression node,
+    ast.ConstantExpression expression,
     ast.InstanceConstant constant,
   ) =>
       InstanceReference(
-        location: node.location!.recordLocation(source),
+        location: expression.location!.recordLocation(_source),
         instanceConstant: _fieldsFromConstant(constant),
-        loadingUnit: loadingUnitForNode(node, _loadingUnits).toString(),
+        loadingUnit: loadingUnitForNode(expression, _loadingUnits).toString(),
       );
 
   InstanceConstant _fieldsFromConstant(ast.InstanceConstant constant) =>
       InstanceConstant(
-          fields: constant.fieldValues.map(
-        (key, value) => MapEntry(
-          key.asField.name.text,
-          evaluateConstant(value),
+        fields: constant.fieldValues.map(
+          (key, value) => MapEntry(
+            key.asField.name.text,
+            evaluateConstant(value),
+          ),
         ),
-      ));
+      );
 }