require 'thread'

##
# A container for data structures that hold weakly on their
# contents.

module Weak

  ##
  # Weak::RefError is raised when the system is unable to
  # dereference an object_id.

  class RefError < RuntimeError; end

  ##
  # A key for Weak::IdentityKeyHash that holds weakly on its
  # contents.

  class Key

    ##
    # The object being weakly held by this key.

    attr_reader :internal_id

    ##
    # Map from an object to a list of its references

    ID_MAP = {}

    ##
    # Map from a reference to its object

    ID_REV_MAP = {}

    ##
    # Creates a finalizer for the object referenced by
    # +internal_id+ that removes that object from +hash+.

    def self.make_finalizer(internal_id, hash_id)
      return proc do |id|
        Thread.exclusive do
          rids = ID_MAP[id]
          if rids then
            for rid in rids
              ID_REV_MAP.delete rid
            end

            ID_MAP.delete id
          end

          rid = ID_REV_MAP[id]

          if rid then
            ID_REV_MAP.delete id
            ID_MAP[rid].delete id
            ID_MAP.delete rid if ID_MAP[rid].empty?
          end
        end

        unless hash_id.nil? then
          hash = ObjectSpace._id2ref hash_id
          hash.delete_if do |k, v|
            k.internal_id == internal_id rescue true
          end unless hash.nil?
        end
      end
    end 

    ##
    # Create a new WeakKey from +orig+ that will remove itself
    # from +hash+when it is no longer referenced.

    def initialize(orig, hash)
      @internal_id = orig.object_id

      ObjectSpace.define_finalizer(orig, self.class.make_finalizer(@internal_id, hash.object_id))
      ObjectSpace.define_finalizer(self, self.class.make_finalizer(@internal_id, nil))

      Thread.exclusive do
        ID_MAP[@internal_id] = [] unless ID_MAP[@internal_id]
      end

      ID_MAP[@internal_id].push self.object_id
      ID_REV_MAP[self.object_id] = @internal_id
    end

    ##
    # A WeakKey's hash is the held object's object_id.

    def hash
      @internal_id
    end

    ##
    # A WeakKey is equal to +other+ if +other+ has the same
    # internal_id or object_id.

    def eql?(other)
      if other.respond_to? :internal_id then
        return other.internal_id == @internal_id
      else
        return other.object_id == @internal_id
      end

      return false
    end

    ##
    # Creates a string that represents the internal state of this
    # object.

    def inspect
      "#<#{self.class}:0x%x @internal_id=0x%x>" % [object_id, @internal_id]
    end

    ##
    # Retrieves an object from its id.

    def get_obj
      unless ID_REV_MAP[self.object_id] == @internal_id then
        raise RefError, "Illegal Reference - probably recycled", caller(2)
      end

      begin
        return ObjectSpace._id2ref(@internal_id)
      rescue RangeError
        raise RefError, "Illegal Reference - probably recycled", caller(2)
      end
    end

    ##
    # Retrives +obj+'s WeakKey reference.

    def self.by_obj(obj)
      ref = nil

      Thread.exclusive do
        rids = ID_MAP[obj.object_id]

        if rids.nil? or rids.empty? then
          raise RefError, "Reference does not exist"
        end

        for rid in rids do
          begin
            ref = ObjectSpace._id2ref rid
          rescue RangeError
            next
          end

          return ref unless ref.nil?
        end

        raise RefError, "Reference does not exist"
      end
    end

  end # class Weak::Key

  ##
  # Creates a Hash that holds weakly on its keys.

  class IdentityKeyHash < Hash

    ##
    # Retrieves +key+ from the hash, if it still exists.

    def [](key)
      return super(Weak::Key.by_obj(key))
    end

    ##
    # Sets +key+ of the hash to +value+.  +value+ is weakly held.

    def []=(key, value)
      ref_key = nil

      unless self.has_key? key then
        ref_key = Weak::Key.new key, self
      else
        ref_key = Weak::Key.by_id key.object_id
      end

      super ref_key, value
    end

    ##
    # Hash's #each

    alias internal_each each

    ##
    # Iterates over the hash, retrieving the weakly referenced
    # objects.

    def each
      key = nil

      super do |ref_key, value|
        begin
          key = ref_key.get_obj
        rescue RefError
          next
        end

        yield key, value
      end
    end

  end # class Weak::IdentityKeyHash

end # module Weak