Refactor DescendantsTracker to leverage native Class#descendants on Ruby 3.1

activesupport/CHANGELOG.md

@@ -1,3 +1,14 @@
+*   `ActiveSupport::DescendantsTracker` now mostly delegate to `Class#descendants` on Ruby 3.1
+
+    Ruby now provides a fast `Class#descendants` making `ActiveSupport::DescendantsTracker` mostly useless.
+
+    As a result the following methods are deprecated:
+
+      - `ActiveSupport::DescendantsTracker.direct_descendants`
+      - `ActiveSupport::DescendantsTracker#direct_descendants`
+
+    *Jean Boussier*
+
 *   Fix the `Digest::UUID.uuid_from_hash` behavior for namespace IDs that are different from the ones defined on `Digest::UUID`.
 
     The new behavior will be enabled by setting the

activesupport/lib/active_support.rb

@@ -116,6 +116,8 @@ module ActiveSupport
   def self.current_attributes_use_thread_variables=(value)
     CurrentAttributes._use_thread_variables = value
   end
+
+  @has_native_class_descendants = Class.method_defined?(:descendants) # RUBY_VERSION >= "3.1"
 end
 
 autoload :I18n, "active_support/i18n"

activesupport/lib/active_support/callbacks.rb

@@ -608,7 +608,7 @@ module ActiveSupport
         # This is used internally to append, prepend and skip callbacks to the
         # CallbackChain.
         def __update_callbacks(name) # :nodoc:
-          ([self] + ActiveSupport::DescendantsTracker.descendants(self)).reverse_each do |target|
+          ([self] + self.descendants).reverse_each do |target|
             chain = target.get_callbacks name
             yield target, chain.dup
           end
@@ -732,7 +732,7 @@ module ActiveSupport
         def reset_callbacks(name)
           callbacks = get_callbacks name
 
-          ActiveSupport::DescendantsTracker.descendants(self).each do |target|
+          self.descendants.each do |target|
             chain = target.get_callbacks(name).dup
             callbacks.each { |c| chain.delete(c) }
             target.set_callbacks name, chain
@@ -825,7 +825,7 @@ module ActiveSupport
           names.each do |name|
             name = name.to_sym
 
-            ([self] + ActiveSupport::DescendantsTracker.descendants(self)).each do |target|
+            ([self] + self.descendants).each do |target|
               target.set_callbacks name, CallbackChain.new(name, options)
             end
 

activesupport/lib/active_support/core_ext/class/subclasses.rb

@@ -18,7 +18,7 @@ class Class
     ObjectSpace.each_object(singleton_class).reject do |k|
       k.singleton_class? || k == self
     end
-  end unless method_defined?(:descendants) # RUBY_VERSION >= "3.1"
+  end unless ActiveSupport.instance_variable_get(:@has_native_class_descendants) # RUBY_VERSION >= "3.1"
 
   # Returns an array with the direct children of +self+.
   #

activesupport/lib/active_support/descendants_tracker.rb

@@ -6,108 +6,160 @@ module ActiveSupport
   # This module provides an internal implementation to track descendants
   # which is faster than iterating through ObjectSpace.
   module DescendantsTracker
-    @@direct_descendants = {}
-
     class << self
       def direct_descendants(klass)
-        descendants = @@direct_descendants[klass]
-        descendants ? descendants.to_a : []
+        ActiveSupport::Deprecation.warn(<<~MSG)
+          ActiveSupport::DescendantsTracker.direct_descendants is deprecated and will be removed in Rails 7.1.
+          Use ActiveSupport::DescendantsTracker.subclasses instead.
+        MSG
+        subclasses(klass)
       end
-      alias_method :subclasses, :direct_descendants
+    end
 
-      def descendants(klass)
-        arr = []
-        accumulate_descendants(klass, arr)
-        arr
-      end
-
-      def clear(only: nil)
-        if only.nil?
-          @@direct_descendants.clear
-          return
+    if ActiveSupport.instance_variable_get(:@has_native_class_descendants) # RUBY_VERSION >= "3.1"
+      class << self
+        def subclasses(klass)
+          klass.subclasses
         end
 
-        @@direct_descendants.each do |klass, direct_descendants_of_klass|
-          if only.member?(klass)
-            @@direct_descendants.delete(klass)
-          else
-            direct_descendants_of_klass.reject! do |direct_descendant_of_class|
-              only.member?(direct_descendant_of_class)
+        def descendants(klass)
+          klass.descendants
+        end
+
+        def clear(only: nil) # :nodoc:
+          # noop
+        end
+
+        def native? # :nodoc:
+          true
+        end
+      end
+
+      def subclasses
+        descendants.select { |descendant| descendant.superclass == self }
+      end
+
+      def direct_descendants
+        ActiveSupport::Deprecation.warn(<<~MSG)
+          ActiveSupport::DescendantsTracker#direct_descendants is deprecated and will be removed in Rails 7.1.
+          Use #subclasses instead.
+        MSG
+        subclasses
+      end
+    else
+      @@direct_descendants = {}
+
+      class << self
+        def subclasses(klass)
+          descendants = @@direct_descendants[klass]
+          descendants ? descendants.to_a : []
+        end
+
+        def descendants(klass)
+          arr = []
+          accumulate_descendants(klass, arr)
+          arr
+        end
+
+        def clear(only: nil) # :nodoc:
+          if only.nil?
+            @@direct_descendants.clear
+            return
+          end
+
+          @@direct_descendants.each do |klass, direct_descendants_of_klass|
+            if only.member?(klass)
+              @@direct_descendants.delete(klass)
+            else
+              direct_descendants_of_klass.reject! do |direct_descendant_of_class|
+                only.member?(direct_descendant_of_class)
+              end
             end
           end
         end
-      end
 
-      # This is the only method that is not thread safe, but is only ever called
-      # during the eager loading phase.
-      def store_inherited(klass, descendant)
-        (@@direct_descendants[klass] ||= DescendantsArray.new) << descendant
-      end
-
-      private
-        def accumulate_descendants(klass, acc)
-          if direct_descendants = @@direct_descendants[klass]
-            direct_descendants.each do |direct_descendant|
-              acc << direct_descendant
-              accumulate_descendants(direct_descendant, acc)
-            end
-          end
-        end
-    end
-
-    def inherited(base)
-      DescendantsTracker.store_inherited(self, base)
-      super
-    end
-
-    def direct_descendants
-      DescendantsTracker.direct_descendants(self)
-    end
-    alias_method :subclasses, :direct_descendants
-
-    def descendants
-      DescendantsTracker.descendants(self)
-    end
-
-    # DescendantsArray is an array that contains weak references to classes.
-    class DescendantsArray # :nodoc:
-      include Enumerable
-
-      def initialize
-        @refs = []
-      end
-
-      def initialize_copy(orig)
-        @refs = @refs.dup
-      end
-
-      def <<(klass)
-        @refs << WeakRef.new(klass)
-      end
-
-      def each
-        @refs.reject! do |ref|
-          yield ref.__getobj__
+        def native? # :nodoc:
           false
-        rescue WeakRef::RefError
-          true
         end
-        self
+
+        # This is the only method that is not thread safe, but is only ever called
+        # during the eager loading phase.
+        def store_inherited(klass, descendant)
+          (@@direct_descendants[klass] ||= DescendantsArray.new) << descendant
+        end
+
+        private
+          def accumulate_descendants(klass, acc)
+            if direct_descendants = @@direct_descendants[klass]
+              direct_descendants.each do |direct_descendant|
+                acc << direct_descendant
+                accumulate_descendants(direct_descendant, acc)
+              end
+            end
+          end
       end
 
-      def refs_size
-        @refs.size
+      def inherited(base)
+        DescendantsTracker.store_inherited(self, base)
+        super
       end
 
-      def cleanup!
-        @refs.delete_if { |ref| !ref.weakref_alive? }
+      def direct_descendants
+        ActiveSupport::Deprecation.warn(<<~MSG)
+          ActiveSupport::DescendantsTracker#direct_descendants is deprecated and will be removed in Rails 7.1.
+          Use #subclasses instead.
+        MSG
+        DescendantsTracker.subclasses(self)
       end
 
-      def reject!
-        @refs.reject! do |ref|
-          yield ref.__getobj__
-        rescue WeakRef::RefError
-          true
+      def subclasses
+        DescendantsTracker.subclasses(self)
+      end
+
+      def descendants
+        DescendantsTracker.descendants(self)
+      end
+
+      # DescendantsArray is an array that contains weak references to classes.
+      class DescendantsArray # :nodoc:
+        include Enumerable
+
+        def initialize
+          @refs = []
+        end
+
+        def initialize_copy(orig)
+          @refs = @refs.dup
+        end
+
+        def <<(klass)
+          @refs << WeakRef.new(klass)
+        end
+
+        def each
+          @refs.reject! do |ref|
+            yield ref.__getobj__
+            false
+          rescue WeakRef::RefError
+            true
+          end
+          self
+        end
+
+        def refs_size
+          @refs.size
+        end
+
+        def cleanup!
+          @refs.delete_if { |ref| !ref.weakref_alive? }
+        end
+
+        def reject!
+          @refs.reject! do |ref|
+            yield ref.__getobj__
+          rescue WeakRef::RefError
+            true
+          end
         end
       end
     end

activesupport/test/descendants_tracker_test.rb

@@ -6,8 +6,10 @@ require "active_support/descendants_tracker"
 
 class DescendantsTrackerTest < ActiveSupport::TestCase
   setup do
-    @original_state = ActiveSupport::DescendantsTracker.class_eval("@@direct_descendants").dup
-    @original_state.each { |k, v| @original_state[k] = v.dup }
+    if ActiveSupport::DescendantsTracker.class_variable_defined?(:@@direct_descendants)
+      @original_state = ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants).dup
+      @original_state.each { |k, v| @original_state[k] = v.dup }
+    end
 
     ActiveSupport::DescendantsTracker.clear
     eval <<~RUBY
@@ -30,10 +32,14 @@ class DescendantsTrackerTest < ActiveSupport::TestCase
   end
 
   teardown do
-    ActiveSupport::DescendantsTracker.class_eval("@@direct_descendants").replace(@original_state)
+    if ActiveSupport::DescendantsTracker.class_variable_defined?(:@@direct_descendants)
+      ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants).replace(@original_state)
+    end
 
     %i(Parent Child1 Child2 Grandchild1 Grandchild2).each do |name|
-      DescendantsTrackerTest.send(:remove_const, name)
+      if DescendantsTrackerTest.const_defined?(name)
+        DescendantsTrackerTest.send(:remove_const, name)
+      end
     end
   end
 
@@ -65,30 +71,42 @@ class DescendantsTrackerTest < ActiveSupport::TestCase
   end
 
   test ".direct_descendants" do
-    assert_equal_sets [Child1, Child2], Parent.direct_descendants
-    assert_equal_sets [Grandchild1, Grandchild2], Child1.direct_descendants
-    assert_equal_sets [], Child2.direct_descendants
+    assert_deprecated do
+      assert_equal_sets [Child1, Child2], Parent.direct_descendants
+    end
+
+    assert_deprecated do
+      assert_equal_sets [Grandchild1, Grandchild2], Child1.direct_descendants
+    end
+
+    assert_deprecated do
+      assert_equal_sets [], Child2.direct_descendants
+    end
   end
 
   test ".subclasses" do
     [Parent, Child1, Child2].each do |klass|
-      assert_equal klass.direct_descendants, klass.subclasses
+      assert_equal assert_deprecated { klass.direct_descendants }, klass.subclasses
     end
   end
 
   test ".clear deletes all state" do
     ActiveSupport::DescendantsTracker.clear
-    assert_empty ActiveSupport::DescendantsTracker.class_eval("@@direct_descendants")
+    if ActiveSupport::DescendantsTracker.class_variable_defined?(:@@direct_descendants)
+      assert_empty ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants)
+    end
   end
 
   test ".clear(only) deletes the given classes only" do
+    skip "Irrelevant for native Class#descendants" if ActiveSupport::DescendantsTracker.native?
+
     ActiveSupport::DescendantsTracker.clear(only: Set[Child2, Grandchild1])
 
     assert_equal_sets [Child1, Grandchild2], Parent.descendants
     assert_equal_sets [Grandchild2], Child1.descendants
 
-    assert_equal_sets [Child1], Parent.direct_descendants
-    assert_equal_sets [Grandchild2], Child1.direct_descendants
+    assert_equal_sets [Child1], assert_deprecated { Parent.direct_descendants }
+    assert_equal_sets [Grandchild2], assert_deprecated { Child1.direct_descendants }
   end
 
   private