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,14 +6,54 @@ module ActiveSupport
   # This module provides an internal implementation to track descendants
   # which is faster than iterating through ObjectSpace.
   module DescendantsTracker
+    class << self
+      def direct_descendants(klass)
+        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
+    end
+
+    if ActiveSupport.instance_variable_get(:@has_native_class_descendants) # RUBY_VERSION >= "3.1"
+      class << self
+        def subclasses(klass)
+          klass.subclasses
+        end
+
+        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 direct_descendants(klass)
+        def subclasses(klass)
           descendants = @@direct_descendants[klass]
           descendants ? descendants.to_a : []
         end
-      alias_method :subclasses, :direct_descendants
 
         def descendants(klass)
           arr = []
@@ -21,7 +61,7 @@ module ActiveSupport
           arr
         end
 
-      def clear(only: nil)
+        def clear(only: nil) # :nodoc:
           if only.nil?
             @@direct_descendants.clear
             return
@@ -38,6 +78,10 @@ module ActiveSupport
           end
         end
 
+        def native? # :nodoc:
+          false
+        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)
@@ -61,9 +105,16 @@ module ActiveSupport
       end
 
       def direct_descendants
-      DescendantsTracker.direct_descendants(self)
+        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 subclasses
+        DescendantsTracker.subclasses(self)
       end
-    alias_method :subclasses, :direct_descendants
 
       def descendants
         DescendantsTracker.descendants(self)
@@ -113,3 +164,4 @@ module ActiveSupport
       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
+    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,12 +32,16 @@ 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|
+      if DescendantsTrackerTest.const_defined?(name)
         DescendantsTrackerTest.send(:remove_const, name)
       end
     end
+  end
 
   test ".descendants" do
     assert_equal_sets [Child1, Grandchild1, Grandchild2, Child2], Parent.descendants
@@ -65,30 +71,42 @@ class DescendantsTrackerTest < ActiveSupport::TestCase
   end
 
   test ".direct_descendants" do
+    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