A self-cleaning weak-ref registry to retrieve object instances by their object_id in Ruby

weak_ref_registry.rb

# This allows doing
#
#     Foo.prepend(WeakRefRegistry)
#
#     foo = Foo.new              # => #<Foo:0x1>
#     foo_id = foo.object_id     # => 1234
#     Foo.registry.fetch(foo_id) # => #<Foo:0x1>
#
# I.e. retrieve Foo instances by their object_id,
# without affecting garbage collection.
#
# This is basically a slighty faster alternative to doing
# `ObjectSpace.each_object(Foo).find { |f| f.object_id == foo_id }`

require 'weakref'

module WeakRefRegistry
  def initialize(...)
    super
    self.class.registry.store(object_id, WeakRef.new(self))
    ObjectSpace.define_finalizer(self, self.class.finalizer)
  end

  def self.included(*) = raise("#{self} must be prepended")

  def self.prepended(klass)
    def klass.registry  = @registry ||= Registry.new
    def klass.finalizer = ->(object_id){ registry.delete(object_id) }
  end

  class Registry
    def fetch(...)  = data.fetch(...)
    def store(...)  = data.store(...)
    def delete(...) = data.delete(...)

    private
    
    def data
      @data ||= {}
    end
  end
end