Ruby `OrderedSet`

ordered_set.rb

require 'set'

if defined?(RUBY_VERSION) && RUBY_VERSION.to_f >= 1.9
  class OrderedSet < Set
    def ==(other)
      return false unless other.is_a?(OrderedSet) && size == other.size
      enum = other.to_enum
      all? { |x| x == enum.next }
    end
    
    def inspect
      super.sub!(/^#</, '#<Ordered')
    end
    
    alias_method :to_ary, :to_a
    
    def to_set
      Set.new(self)
    end
  end
else
  class OrderedSet < Set
    attr_reader :array, :set
    protected :array, :set
  
    def initialize(enumerable = nil)
      replace(enumerable)
    end
  
    def initialize_copy(original)
      super
      replace(original)
    end
  
    def ==(other)
      other.is_a?(OrderedSet) && @array == other.array
    end
  
    def size
      @array.size
    end
  
    def hash
      @array.hash
    end
  
    def replace(enumerable)
      if enumerable.is_a?(OrderedSet)
        @set = enumerable.set.dup
        @array = enumerable.array.dup
    
      else
        @set = Set.new(enumerable)
        @array = Array(enumerable)
      
        if @array.equal?(enumerable)
          @array = @array.size == @set.size ? @array.dup : @array.uniq
        elsif @array.size != @set.size
          @array.uniq!
        end
      end
    
      self
    end
  
    def include?(object)
      @set.include?(object)
    end
  
    def add?(object)
      if @set.add?(object)
        @array << object
        self
      end
    end
  
    def add(object)
      add?(object)
      self
    end
    alias_method :<<, :add
  
    def delete?(object)
      if @set.delete?(object)
        @array.delete(object)
        self
      end
    end
  
    def delete(object)
      delete?(object)
      self
    end
  
    def clear
      @set.clear
      @array.clear
      self
    end
  
    def each
      return enum_for(__method__) unless block_given?
      @array.each { |o| yield(o) }
      self
    end
  
    def delete_if
      return enum_for(__method__) unless block_given?
      @array.dup.each do |o|
        if yield(o)
          @set.delete(o)
          @array.delete(o)
        end
      end
      self
    end
  
    def reject!
      return enum_for(__method__) unless block_given?
      size = @array.size
      delete_if { |o| yield(o) }
      self if @array.size != size
    end
  
    def keep_if
      return enum_for(__method__) unless block_given?
      delete_if { |o| not yield(o) }
      self
    end
  
    def select!
      return enum_for(__method__) unless block_given?
      size = @array.size
      keep_if { |o| yield(o) }
      self if @array.size != size
    end
  
    def collect!
      replace collect { |o| yield(o) }
      self
    end
    alias_method :map!, :collect!
  
    def flatten
      set = self.class.new
      set.flatten_merge(self)
      set
    end
  
    def flatten!
      replace(flatten) if any? { |o| o.is_a?(Set) }
    end
  
    def to_a
      @array.dup
    end
    alias_method :to_ary, :to_a
  
    def to_set
      @set.dup
    end
  
    methods = Set.instance_methods(false)
    methods -= Enumerable.instance_methods
    methods -= OrderedSet.instance_methods
    methods.each do |method|
      class_eval(<<-EVAL, __FILE__, __LINE__)
        def #{method}(*args, &block)
          @set.#{method}(*args, &block)
        end
      EVAL
    end
  
    protected
      def flatten_merge(set, seen = Set.new)
        set.each do |object|
          if object.is_a?(Set)
            if seen.include?(id = object.object_id)
              raise ArgumentError, 'tried to flatten recursive Set'
            end
            seen.add(id)
            flatten_merge(object, seen)
            seen.delete(id)
          else
            add?(object)
          end
        end
      end
  end
end

if __FILE__ == $PROGRAM_NAME
  set = OrderedSet.new([1,2,2,3])
  set << 4 << 5 << 5 << -1
  p set
end