Optimize subclass/subtype queries

pkg/compiler/lib/src/universe/class_set.dart

Index: pkg/compiler/lib/src/universe/class_set.dart
diff --git a/pkg/compiler/lib/src/universe/class_set.dart b/pkg/compiler/lib/src/universe/class_set.dart
index d895de29edc85d86e2cd0b6fc35277210e082a09..3790a03d097c89faf080a82f901e0d006da94eb8 100644
--- a/pkg/compiler/lib/src/universe/class_set.dart
+++ b/pkg/compiler/lib/src/universe/class_set.dart
@@ -6,6 +6,7 @@ library dart2js.world.class_set;
 
 import 'dart:collection' show
     IterableBase;
+import '../common.dart';
 import '../elements/elements.dart' show
     ClassElement;
 import '../util/enumset.dart' show
@@ -92,12 +93,14 @@ class ClassHierarchyNode {
     return mask;
   }
 
+  final ClassHierarchyNode parentNode;
   final ClassElement cls;
   final EnumSet<Instantiation> _mask =
       new EnumSet<Instantiation>.fromValues(
           const <Instantiation>[Instantiation.UNINSTANTIATED]);
 
   ClassElement _leastUpperInstantiatedSubclass;
+  int _instantiatedSubclassCount = 0;
 
   /// `true` if [cls] has been directly instantiated.
   ///
@@ -111,14 +114,23 @@ class ClassHierarchyNode {
 
   void set isDirectlyInstantiated(bool value) {
     if (value != isDirectlyInstantiated) {
+      ClassHierarchyNode parent = parentNode;
       if (value) {
         _mask.remove(Instantiation.UNINSTANTIATED);
         _mask.add(Instantiation.DIRECTLY_INSTANTIATED);
+        while (parent != null) {
+          parent._updateInstantiatedSubclassCount(1);
+          parent = parent.parentNode;
+        }
       } else {
         _mask.remove(Instantiation.DIRECTLY_INSTANTIATED);
         if (_mask.isEmpty) {
           _mask.add(Instantiation.UNINSTANTIATED);
         }
+        while (parent != null) {
+          parent._updateInstantiatedSubclassCount(-1);
+          parent = parent.parentNode;
+        }
       }
     }
   }
@@ -131,12 +143,18 @@ class ClassHierarchyNode {
   ///   class C extends B {}
   ///   main() => [new B(), new C()];
   ///
-  bool get isIndirectlyInstantiated =>
-      _mask.contains(Instantiation.INDIRECTLY_INSTANTIATED);
-
-  void set isIndirectlyInstantiated(bool value) {
-    if (value != isIndirectlyInstantiated) {
-      if (value) {
+  bool get isIndirectlyInstantiated => _instantiatedSubclassCount > 0;
+
+  /// The number of strict subclasses that are directly or indirectly
+  /// instantiated.
+  int get instantiatedSubclassCount => _instantiatedSubclassCount;
+
+  void _updateInstantiatedSubclassCount(int change) {
+    bool before = isIndirectlyInstantiated;
+    _instantiatedSubclassCount += change;
+    bool after = isIndirectlyInstantiated;
+    if (before != after) {
+      if (after) {
         _mask.remove(Instantiation.UNINSTANTIATED);
         _mask.add(Instantiation.INDIRECTLY_INSTANTIATED);
       } else {
@@ -151,7 +169,11 @@ class ClassHierarchyNode {
   /// The nodes for the direct subclasses of [cls].
   Link<ClassHierarchyNode> _directSubclasses = const Link<ClassHierarchyNode>();
 
-  ClassHierarchyNode(this.cls);
+  ClassHierarchyNode(this.parentNode, this.cls) {
+    if (parentNode != null) {
+      parentNode.addDirectSubclass(this);
+    }
+  }
 
   /// Adds [subclass] as a direct subclass of [cls].
   void addDirectSubclass(ClassHierarchyNode subclass) {
@@ -177,24 +199,6 @@ class ClassHierarchyNode {
 
   /// Returns an [Iterable] of the subclasses of [cls] possibly including [cls].
   ///
-  /// The directly instantiated, indirectly instantiated and uninstantiated
-  /// subclasses of [cls] are returned if [includeDirectlyInstantiated],
-  /// [includeIndirectlyInstantiated], and [includeUninstantiated] are `true`,
-  /// respectively. If [strict] is `true`, [cls] itself is _not_ returned.
-  Iterable<ClassElement> subclasses(
-      {bool includeDirectlyInstantiated: true,
-       bool includeIndirectlyInstantiated: true,
-       bool includeUninstantiated: true,
-       bool strict: false}) {
-    EnumSet<Instantiation> mask = createMask(
-        includeDirectlyInstantiated: includeDirectlyInstantiated,
-        includeIndirectlyInstantiated:includeIndirectlyInstantiated,
-        includeUninstantiated: includeUninstantiated);
-    return subclassesByMask(mask, strict: strict);
-  }
-
-  /// Returns an [Iterable] of the subclasses of [cls] possibly including [cls].
-  ///
   /// Subclasses are included if their instantiation properties intersect with
   /// their corresponding [Instantiation] values in [mask]. If [strict] is
   /// `true`, [cls] itself is _not_ returned.
@@ -205,6 +209,65 @@ class ClassHierarchyNode {
         this, mask, includeRoot: !strict);
   }
 
+  /// Applies [predicate] to each subclass of [cls] matching the criteria
+  /// specified by [mask] and [strict]. If [predicate] returns `true` on a
+  /// class, visitation is stopped immediately and the function returns `true`.
+  ///
+  /// [predicate] is applied to subclasses if their instantiation properties
+  /// intersect with their corresponding [Instantiation] values in [mask]. If
+  /// [strict] is `true`, [predicate] is _not_ called on [cls] itself.
+  bool anySubclass(
+      bool predicate(ClassElement cls),
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+
+    ForEach wrapper(ClassElement cls) {
+      return predicate(cls) ? ForEach.STOP : ForEach.CONTINUE;
+    }
+    return forEachSubclass(wrapper, mask, strict: strict) == ForEach.STOP;
+  }
+
+  /// Applies [f] to each subclass of [cls] matching the criteria specified by
+  /// [mask] and [strict].
+  ///
+  /// [f] is a applied to subclasses if their instantiation properties intersect
+  /// with their corresponding [Instantiation] values in [mask]. If [strict] is
+  /// `true`, [f] is _not_ called on [cls] itself.
+  ///
+  /// The visitation of subclasses can be cut short by the return value of [f].
+  /// If [ForEach.STOP] is returned, no further classes are visited and the
+  /// function stops immediately. If [ForEach.SKIP_SUBCLASSES] is returned, the
+  /// subclasses of the last visited class are skipped, but visitation
+  /// continues. The return value of the function is either [ForEach.STOP], if
+  /// visitation was stopped, or [ForEach.CONTINUE] if visitation continued to
+  /// the end.
+  ForEach forEachSubclass(
+      ForEachFunction f,
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+    ForEach forEach;
+    if (!strict && mask.intersects(_mask)) {
+      forEach = f(cls);
+    }
+    // Interpret `forEach == null` as `forEach == ForEach.CONTINUE`.
+    forEach ??= ForEach.CONTINUE;
+
+    if (forEach == ForEach.CONTINUE) {
+      if (mask.contains(Instantiation.UNINSTANTIATED) || isInstantiated) {
+        for (ClassHierarchyNode subclass in _directSubclasses) {
+          ForEach subForEach = subclass.forEachSubclass(f, mask);
+          if (subForEach == ForEach.STOP) {
+            return subForEach;
+          }
+        }
+      }
+    }
+    if (forEach == ForEach.STOP) {
+      return forEach;
+    }
+    return ForEach.CONTINUE;
+  }
+
   /// Returns the most specific subclass of [cls] (including [cls]) that is
   /// directly instantiated or a superclass of all directly instantiated
   /// subclasses. If [cls] is not instantiated, `null` is returned.
@@ -275,7 +338,8 @@ class ClassHierarchyNode {
         if (instantiatedOnly && !child.isInstantiated) {
           continue;
         }
-        if (withRespectTo != null && !child.subclasses().any(isRelatedTo)) {
+        if (withRespectTo != null &&
+            !child.anySubclass(isRelatedTo, ClassHierarchyNode.ALL)) {
           continue;
         }
         if (needsComma) {
@@ -367,22 +431,31 @@ class ClassSet {
 
   ClassElement get cls => node.cls;
 
-  /// Returns an [Iterable] of the subclasses of [cls] possibly including [cls].
-  ///
-  /// The directly instantiated, indirectly instantiated and uninstantiated
-  /// subclasses of [cls] are returned if [includeDirectlyInstantiated],
-  /// [includeIndirectlyInstantiated], and [includeUninstantiated] are `true`,
-  /// respectively. If [strict] is `true`, [cls] itself is _not_ returned.
-  Iterable<ClassElement> subclasses(
-      {bool includeDirectlyInstantiated: true,
-       bool includeIndirectlyInstantiated: true,
-       bool includeUninstantiated: true,
-       bool strict: false}) {
-    EnumSet<Instantiation> mask = ClassHierarchyNode.createMask(
-        includeDirectlyInstantiated: includeDirectlyInstantiated,
-        includeIndirectlyInstantiated:includeIndirectlyInstantiated,
-        includeUninstantiated: includeUninstantiated);
-    return subclassesByMask(mask, strict: strict);
+  /// Returns the number of directly instantiated subtypes of [cls].
+  int get instantiatedSubtypeCount {
+    int count = node.instantiatedSubclassCount;
+    if (_directSubtypes != null) {
+      for (ClassHierarchyNode subtypeNode in _directSubtypes) {
+        if (subtypeNode.isDirectlyInstantiated) {
+          count++;
+        }
+        count += subtypeNode.instantiatedSubclassCount;
+      }
+    }
+    return count;
+  }
+
+  /// Returns `true` if all instantiated subtypes of [cls] are subclasses of
+  /// [cls].
+  bool get hasOnlyInstantiatedSubclasses {

[Siggi Cherem (dart-lang), 2016/02/09 23:40:12]

I was looking into something else and run into this part of the change. I think
I understand how this works now, but I wanted to confirm/double check with you.

In particular, I was thinking of this example:

  class B extends A {}
  class C extends B {}
  class D implements B {}
  main() => print([new A(), new B(), new C(), new D()]);

I expect hasOnlyInstantiatedSubclasses for `A` to be false. I believe that's the
case here because the 'directSubtypes' includes subtypes of subclasses too, so
in this case `D` is in the subtypes list of `A`'s class-set, and because of that
we'll say false here.

Is that right?

[Johnni Winther, 2016/02/10 09:32:08]

Yes. The cause is that implements is transitive. The
[hasOnlyInstantiatedSubclasses] is checked to know if `subtype of A` means
`extends A`, and because of `D` we say no. We should, though, rename
[directSubtypes] to something else.

+    if (_directSubtypes != null) {
+      for (ClassHierarchyNode subtypeNode in _directSubtypes) {
+        if (subtypeNode.isInstantiated) {
+          return false;
+        }
+      }
+    }
+    return true;
   }
 
   /// Returns an [Iterable] of the subclasses of [cls] possibly including [cls].
@@ -414,7 +487,6 @@ class ClassSet {
     return subtypesByMask(mask, strict: strict);
   }
 
-
   /// Returns an [Iterable] of the subtypes of [cls] possibly including [cls].
   ///
   /// Subtypes are included if their instantiation properties intersect with
@@ -434,6 +506,91 @@ class ClassSet {
         includeRoot: !strict);
   }
 
+  /// Applies [predicate] to each subclass of [cls] matching the criteria
+  /// specified by [mask] and [strict]. If [predicate] returns `true` on a
+  /// class, visitation is stopped immediately and the function returns `true`.
+  ///
+  /// [predicate] is applied to subclasses if their instantiation properties
+  /// intersect with their corresponding [Instantiation] values in [mask]. If
+  /// [strict] is `true`, [predicate] is _not_ called on [cls] itself.
+  bool anySubclass(
+      bool predicate(ClassElement cls),
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+    return node.anySubclass(predicate, mask, strict: strict);
+  }
+
+  /// Applies [f] to each subclass of [cls] matching the criteria specified by
+  /// [mask] and [strict].
+  ///
+  /// [f] is a applied to subclasses if their instantiation properties intersect
+  /// with their corresponding [Instantiation] values in [mask]. If [strict] is
+  /// `true`, [f] is _not_ called on [cls] itself.
+  ///
+  /// The visitation of subclasses can be cut short by the return value of [f].
+  /// If [ForEach.STOP] is returned, no further classes are visited and the
+  /// function stops immediately. If [ForEach.SKIP_SUBCLASSES] is returned, the
+  /// subclasses of the last visited class are skipped, but visitation
+  /// continues. The return value of the function is either [ForEach.STOP], if
+  /// visitation was stopped, or [ForEach.CONTINUE] if visitation continued to
+  /// the end.
+  ForEach forEachSubclass(
+      ForEachFunction f,
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+    return node.forEachSubclass(f, mask, strict: strict);
+  }
+
+  /// Applies [predicate] to each subtype of [cls] matching the criteria
+  /// specified by [mask] and [strict]. If [predicate] returns `true` on a
+  /// class, visitation is stopped immediately and the function returns `true`.
+  ///
+  /// [predicate] is applied to subtypes if their instantiation properties
+  /// intersect with their corresponding [Instantiation] values in [mask]. If
+  /// [strict] is `true`, [predicate] is _not_ called on [cls] itself.
+  bool anySubtype(
+      bool predicate(ClassElement cls),
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+
+    ForEach wrapper(ClassElement cls) {
+      return predicate(cls) ? ForEach.STOP : ForEach.CONTINUE;
+    }
+    return forEachSubtype(wrapper, mask, strict: strict) == ForEach.STOP;
+  }
+
+  /// Applies [f] to each subtype of [cls] matching the criteria specified by
+  /// [mask] and [strict].
+  ///
+  /// [f] is a applied to subtypes if their instantiation properties intersect
+  /// with their corresponding [Instantiation] values in [mask]. If [strict] is
+  /// `true`, [f] is _not_ called on [cls] itself.
+  ///
+  /// The visitation of subtypes can be cut short by the return value of [f].
+  /// If [ForEach.STOP] is returned, no further classes are visited and the
+  /// function stops immediately. If [ForEach.SKIP_SUBCLASSES] is returned, the
+  /// subclasses of the last visited class are skipped, but visitation
+  /// continues. The return value of the function is either [ForEach.STOP], if
+  /// visitation was stopped, or [ForEach.CONTINUE] if visitation continued to
+  /// the end.
+  ForEach forEachSubtype(
+      ForEachFunction f,
+      EnumSet<Instantiation> mask,
+      {bool strict: false}) {
+    ForEach forEach = node.forEachSubclass(f, mask, strict: strict);
+    forEach ??= ForEach.CONTINUE;

[Siggi Cherem (dart-lang), 2016/02/06 00:17:26]

nit: here and elsewhere, I rather keep the variable with a single assignment and
use ?? instead, that makes it more obvious that it is a default value. In other
words:

  var x = something ?? defaultValue;

instead of 

  var x = something;
  x ??= defaultValue;

[Johnni Winther, 2016/02/10 09:32:08]

Will do. Code motion caused this ;-)

+    if (forEach == ForEach.CONTINUE && _directSubtypes != null) {
+      for (ClassHierarchyNode subclass in _directSubtypes) {
+        ForEach subForEach = subclass.forEachSubclass(f, mask);
+        if (subForEach == ForEach.STOP) {
+          return subForEach;
+        }
+      }
+    }
+    assert(forEach != ForEach.SKIP_SUBCLASSES);
+    return forEach;
+  }
+
   /// Adds [subtype] as a subtype of [cls].
   void addSubtype(ClassHierarchyNode subtype) {
     if (node.contains(subtype.cls)) {
@@ -695,3 +852,20 @@ class SubtypesIterator extends Iterator<ClassElement> {
     return false;
   }
 }
+
+/// Enum values returned from the [ForEachFunction] provided to the `forEachX`
+/// functions of [ClassHierarchyNode] and [ClassSet]. The value is used to
+/// control the continued iteration.
+enum ForEach {

[Siggi Cherem (dart-lang), 2016/02/06 00:17:26]

naming nit: could we rename this to something that describes what this is,
rather than where it is used?

Some ideas:
- enum IterationControl { CONTINUE, STOP, SKIP_SUBCLASSES }
- enum IterationStep    { NEXT_CLASS, STOP, NEXT_SUBCLASS }

I like the latter, since the corresponding local variable names in the code
above become more readable. For example:

var nextStep =
    node.forEachSubclass(f, mask, strict: strict) ?? IterationStep.NEXT_CLASS;

if (nextStep == ForEach.NEXT_CLASS && _directSubtypes != null) {
  for (ClassHierarchyNode subclass in _directSubtypes) {
    var subStep = subclass.forEachSubclass(f, mask);
    if (subStep == IterationStep.STOP) return IterationStep.STOP;
  }
}

return nextStep;

[Johnni Winther, 2016/02/10 09:32:08]

If find the latter confusing but like enum class name.

+  /// Iteration continues.
+  CONTINUE,
+  /// Iteration stops immediately.
+  STOP,
+  /// Iteration skips the subclasses of the current class.
+  SKIP_SUBCLASSES,
+}
+
+/// Visiting function used for the `forEachX` functions of [ClassHierarchyNode]
+/// and [ClassSet]. The return value controls the continued iteration. If `null`
+/// is returned, iteration continues to the end.
+typedef ForEach ForEachFunction(ClassElement cls);