Add ActiveSupport::DescendantsTracker.

activesupport/lib/active_support.rb

@@ -51,6 +51,7 @@ module ActiveSupport
     autoload :Concern
     autoload :Configurable
     autoload :Deprecation
+    autoload :DescendantsTracker
     autoload :Gzip
     autoload :Inflector
     autoload :JSON

activesupport/lib/active_support/descendants_tracker.rb

@@ -0,0 +1,36 @@
+require 'active_support/dependencies'
+
+module ActiveSupport
+  # This module provides an internal implementation to track descendants
+  # which is faster than iterating through ObjectSpace.
+  module DescendantsTracker
+    mattr_accessor :descendants
+    @@descendants = Hash.new { |h, k| h[k] = [] }
+
+    def self.clear
+      @@descendants.each do |klass, descendants|
+        if ActiveSupport::Dependencies.autoloaded?(klass)
+          @@descendants.delete(klass)
+        else
+          descendants.reject! { |v| ActiveSupport::Dependencies.autoloaded?(v) }
+        end
+      end
+    end
+
+    def inherited(base)
+      self.direct_descendants << base
+      super
+    end
+
+    def direct_descendants
+      @@descendants[self]
+    end
+
+    def descendants
+      @@descendants[self].inject([]) do |descendants, klass|
+        descendants << klass
+        descendants.concat klass.descendants
+      end
+    end
+  end
+end

activesupport/test/descendants_tracker_test.rb

@@ -0,0 +1,72 @@
+require 'abstract_unit'
+require 'test/unit'
+require 'active_support'
+
+class DescendantsTrackerTest < Test::Unit::TestCase
+  class Parent
+    extend ActiveSupport::DescendantsTracker
+  end
+
+  class Child1 < Parent
+  end
+
+  class Child2 < Parent
+  end
+
+  class Grandchild1 < Child1
+  end
+
+  class Grandchild2 < Child1
+  end
+
+  def test_descendants
+    assert_equal [Child1, Grandchild1, Grandchild2, Child2], Parent.descendants
+    assert_equal [Grandchild1, Grandchild2], Child1.descendants
+    assert_equal [], Child2.descendants
+  end
+
+  def test_direct_descendants
+    assert_equal [Child1, Child2], Parent.direct_descendants
+    assert_equal [Grandchild1, Grandchild2], Child1.direct_descendants
+    assert_equal [], Child2.direct_descendants
+  end
+
+  def test_clear_with_autoloaded_parent_children_and_granchildren
+    mark_as_autoloaded Parent, Child1, Child2, Grandchild1, Grandchild2 do
+      ActiveSupport::DescendantsTracker.clear
+      assert_equal Hash.new, ActiveSupport::DescendantsTracker.descendants
+    end
+  end
+
+  def test_clear_with_autoloaded_children_and_granchildren
+    mark_as_autoloaded Child1, Grandchild1, Grandchild2 do
+      ActiveSupport::DescendantsTracker.clear
+      assert_equal [Child2], Parent.descendants
+      assert_equal [], Child2.descendants
+    end
+  end
+
+  def test_clear_with_autoloaded_granchildren
+    mark_as_autoloaded Grandchild1, Grandchild2 do
+      ActiveSupport::DescendantsTracker.clear
+      assert_equal [Child1, Child2], Parent.descendants
+      assert_equal [], Child1.descendants
+      assert_equal [], Child2.descendants
+    end
+  end
+
+  protected
+
+  def mark_as_autoloaded(*klasses)
+    old_autoloaded = ActiveSupport::Dependencies.autoloaded_constants.dup
+    ActiveSupport::Dependencies.autoloaded_constants = klasses.map(&:name)
+
+    old_descendants = ActiveSupport::DescendantsTracker.descendants.dup
+    old_descendants.each { |k, v| old_descendants[k] = v.dup }
+
+    yield
+  ensure
+    ActiveSupport::Dependencies.autoloaded_constants = old_autoloaded
+    ActiveSupport::DescendantsTracker.descendants = old_descendants
+  end
+end