File: sized_ordered_set.rb

package info (click to toggle)
ruby-test-prof 0.12.2%2Bdfsg-1
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, forky, sid, trixie
  • size: 508 kB
  • sloc: ruby: 4,075; makefile: 4
file content (65 lines) | stat: -rw-r--r-- 1,300 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# frozen_string_literal: true

module TestProf
  module Utils
    # Ordered set with capacity
    class SizedOrderedSet
      unless [].respond_to?(:bsearch_index)
        require "test_prof/ext/array_bsearch_index"
        using ArrayBSearchIndex
      end

      include Enumerable

      def initialize(max_size, sort_by: nil, &block)
        @max_size = max_size
        @comparator =
          if block_given?
            block
          elsif !sort_by.nil?
            ->(x, y) { x[sort_by] >= y[sort_by] }
          else
            ->(x, y) { x >= y }
          end
        @data = []
      end

      def <<(item)
        return if data.size == max_size &&
          comparator.call(data.last, item)

        # Find an index of a smaller element
        index = data.bsearch_index { |x| !comparator.call(x, item) }

        if index.nil?
          data << item
        else
          data.insert(index, item)
        end

        data.pop if data.size > max_size
        data.size
      end

      def each(&block)
        if block_given?
          data.each(&block)
        else
          data.each
        end
      end

      def size
        data.size
      end

      def to_a
        data.dup
      end

      private

      attr_reader :max_size, :data, :comparator
    end
  end
end