#!/usr/bin/env ruby
# A modification of Joel VanderWerf's EnumerableOperator#product to include 
# restrictions ('select' style filters that are applied as the product is
# constructed).
# 
# The argument list consists of a mixture of factors (Enumerable objects)
# and restrictions (procs returning a boolean). 
#
# Factors appear in order; restrictions are sorted by arity, so that a
# restriction of arity n is applied after the nth element in every tuple is
# added.

require 'tools'

class NilClass
  def and
    return true
  end
end

module EnumerableOperator

  class Product
    include Enumerable

    attr_reader :factors, :dim

    def initialize(*factors)
      @factors = factors
      @filters = []
      @factors.select {|i| i.is_a?(Proc)}.each {|fn|
        (@filters[(fn.arity.abs) -1] ||= []) << fn} 
        # ( a kludge to get around the fact that {|x|}.arity = -1 )
      @factors.reject!{|i| i.is_a?(Proc)}
      @dim = @factors.length
    end

    def each tuple = [nil]*@dim, i = 0, &block
      @factors[i].each {|x| 
        tuple[i] = x
        args = (i == 0) ? x : tuple[0..i]
        if (@filters[i].and {|f| f.call(*args)})
          if i == @dim - 1 then
            yield tuple.dup 
					else
            each tuple, i + 1, &block 
          end
        end
      }
      self if i==0
    end

    def size
      @factors.product { |enum| enum.size }
    end
  end

  def product(*factors, &block)
    if block
      Product.new(*factors).each(&block)
    else
      Product.new(*factors)
    end
  end
  alias :tuples :product

  module_function :product
end

if __FILE__ == $0

  include EnumerableOperator

  for i in product(1..10, proc {|x| x < 5}, 
                    5..9, proc{|x,y| x+y > 8},
                    21..30, proc{|x,y,z| x * y < z}
                   )
     p i
  end

  for x,y,z in product 1..40, 
                       1..40, proc {|x,y| x <= y},
                       1..40, proc {|x,y,z| x**2 + y**2 == z**2}
    p [x,y,z]
  end

  # restrictions can appear in any order...                     
  product(1..40, 1..40, 1..40, proc{|x,y| x <= y}, proc{|x,y,z| y <= z},
            proc{|x,y,z| x**2 + y**2 == z**2}).each {|x,y,z| 
              puts "#{x}^2 + #{y}^2 = #{z}^2"}
              
  # but note, for instance, that [y <= z] needs a three argument filter, even
  # though x does not appear
end