# frozen_string_literal: true

module Pbt
  module Arbitrary
    # Generates arrays of values generated by `value_arb` with type `T`.
    class ArrayArbitrary < Arbitrary
      DEFAULT_MAX_SIZE = 10
      private_constant :DEFAULT_MAX_SIZE

      # @param value_arb [Arbitrary] Arbitrary to generate values.
      # @param min_length [Integer] Minimum length of the generated array. Default is 0.
      # @param max_length [Integer] Maximum length of the generated array. Default is 10.
      def initialize(value_arb, min_length = 0, max_length = DEFAULT_MAX_SIZE)
        raise ArgumentError, "min_length must be zero or positive number" if min_length < 0

        @min_length = min_length
        @max_length = max_length
        @value_arb = value_arb
        @length_arb = IntegerArbitrary.new(min_length, max_length)
      end

      # @see Arbitrary#generate
      def generate(rng)
        length = @length_arb.generate(rng)
        length.times.map { @value_arb.generate(rng) }
      end

      # @see Arbitrary#shrink
      def shrink(current)
        return Enumerator.new { |_| } if current.size == @min_length

        Enumerator.new do |y|
          # First, shrink the length of the array and try combinations of elements.
          # But this doesn't try all possible combinations since it'd be too huge and slow.
          @length_arb.shrink(current.size).each do |length|
            if length == 0
              y.yield []
              next
            end
            current.each_cons(length) do |con|
              y.yield con
            end
          end

          # Second, shrink each element of the array.
          current.each_with_index do |item, i|
            @value_arb.shrink(item).each do |val|
              y.yield [*current[...i], val, *current[i + 1..]]
            end
          end
        end
      end
    end
  end
end
