top_n.rb in top_n-1.1.0

- old
+ new

@@ -12,40 +12,33 @@
 # @note
 #   Note that while the number of keys is restricted, the values are not.
 #   This can cause memory issues if many values are added to the same key.
 #   If this is the type of data you are tracking, you may need a different
 #   solution.
-class TopN
+class TopN < Hash
   # The maxinum number of keys which will be tracked.
   # @return [Fixnum] the configured maximum number of keys to be tracked.
-  attr_reader :maxsize
+  attr_reader :maxkeys
 
   # The configured direction.
   # @return [Symbol] either :top or :bottom.
   attr_reader :direction
 
-  # The current number of keys we are tracking.
-  # @return [FixNum] the count, which will be 0 up to :maxsize.
-  attr_reader :size
-
   # The current value of the minimum (:top) or maximum (:bottom) key.
   # @return [Object] the threshold key.
   attr_reader :threshold_key
 
-  # The currently tracked data as one blob.  This is tied to the
-  # current implementation, so its use is not recommended.
-  # @return [Hash] the internal data structure containing the keys and
-  #  values.  The keys to the returned Hash are the tracked keys, and
-  #  the values at those keys is the list of values.
-  attr_reader :data
+  alias_method :super_store, :store
+  alias_method :super_bracket_assign, :[]=
+  alias_method :super_bracket, :[]
 
   ##
   # Create a new TopN object.  Options available:
   #
   # @param [Hash] options the options used to configure the TopN object.
   #
-  # @option options [Fixnum] :maxsize The maximum number of keys to track.
+  # @option options [Fixnum] :maxkeys The maximum number of keys to track.
   #  Must be a positive Fixnum.  Defaults to 100.
   #
   # @option options [Symbol] :direction Configure the direction.
   #  If this is :top, the largest keys will be maintained.  If :bottom,
   #  the smallest keys will be maintained.  Any other value throws an
@@ -56,41 +49,41 @@
   #
   # @example Create with default options
   #   topn = TopN.new
   #
   # @example Create with a maximum size of 10, and track smaller values
-  #   topn = TopN.new(maxsize: 10, direction: :bottom)
+  #   topn = TopN.new(maxkeys: 10, direction: :bottom)
   #
-  def initialize(options = {})
+  def initialize(default = nil, options = {})
     options = {
-      maxsize: 100,
+      maxkeys: 100,
       direction: :top,
     }.merge(options)
 
     options.keys.each do |opt|
-      unless [:maxsize, :direction].include?opt
+      unless [:maxkeys, :direction].include?opt
         raise ArgumentError.new("invalid option #{opt}")
       end
     end
 
-    @maxsize = options[:maxsize]
+    @maxkeys = options[:maxkeys]
     @direction = options[:direction]
-    @data = {}
-    @size = 0
     @threshold_key = nil
 
     unless [:top, :bottom].include?(@direction)
       raise ArgumentError.new("direction must be :top or :bottom")
     end
 
-    unless @maxsize.is_a?Fixnum
-      raise ArgumentError.new("maxsize must be a Fixnum")
+    unless @maxkeys.is_a?Fixnum
+      raise ArgumentError.new("maxkeys must be a Fixnum")
     end
 
-    if @maxsize <= 0
-      raise ArgumentError.new("maxsize must be >= 1")
+    if @maxkeys <= 0
+      raise ArgumentError.new("maxkeys must be >= 1")
     end
+
+    super(default)
   end
 
   ##
   # Add a key, value pair.
   #
@@ -103,82 +96,62 @@
   # of values at that key.
   #
   # If an existing (key, value) is permitted, and will result in the list of
   # values at that key having the same value multiple times.
   #
-  # @return [Array] if the value was added to the key's list.
-  #
-  # @return [nil] if the value was not added because the key is too small or
-  # large to be tracked.
-  def add(key, value)
+  # @return [Object] the value passed in.  This will be returned even if
+  # the value is not added because there are too many keys already present.
+  def store(key, value)
     if @direction == :top
       add_top(key, value)
     else
       add_bottom(key, value)
     end
   end
 
-  ##
-  # Find and return the list of values for a key.
-  #
-  # @return [Array<Object>] the list of values for 'key'.
-  #
-  # @return [nil] if the key does not exist.
-  def find(key)
-    @data[key]
+  # Behave like #store, with the same semantics.
+  def []=(key, value)
+    store(key, value)
   end
 
-  ##
-  # Return the list of currently tracked keys.
-  #
-  # @return [Array<Object>] the list of values for this key.
-  #  Order is not guaranteed to match the oder which they were added.
-  def keys
-    @data.keys
-  end
-
   private
 
   ##
   # Add a (key, value) when the direction is :top.
   def add_top(key, value)
     @threshold_key ||= key
 
-    if @data.has_key?key
-      @data[key] << value
+    if has_key?key
+      fetch(key) << value
     else
-      if @size >= @maxsize
-        return nil if key < @threshold_key
-        @data.delete(@threshold_key)
-        @size -= 1
-        @threshold_key = @data.keys.min
+      if size >= @maxkeys
+        return value if key < @threshold_key
+        delete(@threshold_key)
+        @threshold_key = keys.min
       end
-      @data[key] = [ value ]
-      @size += 1
+      super_store(key, [ value ])
       @threshold_key = key if key < @threshold_key
     end
 
-    @data[key]
+    value
   end
 
   ##
   # Add a (key, value) when the direction is :bottom.
   def add_bottom(key, value)
     @threshold_key ||= key
 
-    if @data.has_key?key
-      @data[key] << value
+    if has_key?key
+      fetch(key) << value
     else
-      if @size >= @maxsize
-        return nil if key > @threshold_key
-        @data.delete(@threshold_key)
-        @size -= 1
-        @threshold_key = @data.keys.max
+      if size >= @maxkeys
+        return value if key > @threshold_key
+        delete(@threshold_key)
+        @threshold_key = keys.max
       end
-      @data[key] = [ value ]
-      @size += 1
+      super_store(key, [ value ])
       @threshold_key = key if key > @threshold_key
     end
 
-    @data[key]
+    value
   end
 end