class Array
  def nilify; empty? ? nil : self; end
  def unify; (0..1) === size ? first : self; end
end

module Memo
  module ClassMethods
    def memoed_methods
      @memoed_methods ||= {}
    end

    # provide a legal ivar name from a method name. instance variables
    # can't have ? ! and other punctuation. Which isn't handled. Obviously.
    def ivar_from( maybe_meth )
      "@#{maybe_meth.to_s.tr '?!','pi'}".intern
    end

    # store the original method, replace it with a method
    # that memoizes the result.
    def memo( meth )
      unbound_previous_method = instance_method meth
      raise "can't memo #{meth} with arity #{unbound_previous_method.arity}" if unbound_previous_method.arity != 0

      memoed_methods[meth] = unbound_previous_method
      ivar = ivar_from meth

      define_method meth do |*args, &blk|
        if instance_variable_defined? ivar
          instance_variable_get ivar
        else
          # bind the saved method to this instance, call the result ...
          to_memo = unbound_previous_method.bind( self ).call( *args, &blk )
          # ... memo it and return value
          instance_variable_set ivar, to_memo
        end
      end
    end
  end

  # hook in class methods on include
  def self.included( other_module )
    other_module.extend ClassMethods
  end

  # reset some or all memoized variables
  # return cleared values
  def clear_memos( *requested_meths )
    (requested_meths.nilify || self.class.memoed_methods.keys).map do |meth|
      if instance_variable_defined? ivar = self.class.ivar_from(meth)
        remove_instance_variable ivar
      end
    end.unify
  end
end

class YourFunkyChunkyCode
  include Memo

  # for 2.1
  memo def expensive
    # do various things
  end

  # for 2.0
  def even_more_expensive
    # do various more things
  end

  memo :even_more_expensive
end

yfcc = YourFunkyChunkyCode.new
yfcc.expensive

# the currently memo-ed values, and other stuff
ivars = yfcc.instance_variables

# and all the possible memo-ed values (some of which don't exist yet)
memo_ivars = YourFunkyChunkyCode.memoed_methods.keys.map{|meth_name| YourFunkyChunkyCode.ivar_from meth_name}

# only currently memo-ed values
ivars & memo_ivars

require 'memo.rb'
require 'faker'

describe Array do
  def random_values( range )
    rand(range).times.map{|i| [nil,rand,Faker::Name.name,Object.new].sample}
  end

  describe '#unify' do
    it 'nil for 0 elements' do
      [].unify.should == nil
    end

    it 'atom for 1 element' do
      value = rand
      [value].unify.should == value
    end

    it 'array for >=2 elements' do
      random_values(2..15).unify.should be_a(Array)
    end
  end

  describe '#nilify' do
    it 'nil for 0 elements' do
      [].nilify.should == nil
    end

    it 'array for >=1 elements' do
      random_values(1..15).nilify.should be_a(Array)
    end
  end
end

describe Memo::ClassMethods do
  let (:subject) { Object.new.extend(Memo::ClassMethods) }

  describe '#ivar_from' do
    it 'ivars a normal method to symbol' do
      name = Faker::Name.name
      subject.ivar_from(name).should == "@#{name}".to_sym
    end

    it 'ivars a ? method' do
      name = "are_you_mad?"
      subject.ivar_from(name).should == "@are_you_madp".to_sym
    end

    it 'ivars a ! method' do
      name = "jetais_perdu!"
      subject.ivar_from(name).should == "@jetais_perdui".to_sym
    end
  end
end

describe Memo do
  describe '.memo' do
    it 'raises on non-zero arity' do
      class_def = -> do
        Class.new do
          include Memo

          memo def calc( x )
          end
        end
      end
      class_def.should raise_error(/with arity 1/)
    end

    it 'stores previous method' do
      kl = Class.new do
        include Memo
        memo def calc
          rand
        end
      end

      name, body = kl.memoed_methods.first

      name.should == :calc

      body.name.should == :calc
      body.should be_a(UnboundMethod)
    end

    it 'calls previous method once' do
      inst = Class.new do
        attr_reader :call_count
        include Memo
        memo def calc
          @call_count ||= 0
          @call_count += 1
          rand
        end
      end.new

      3.times{inst.calc}
      inst.call_count.should == 1
    end

    it 'memos value' do
      inst = Class.new do
        include Memo
        memo def calc; rand; end
      end.new

      # value should be the same for all calls, even though it's rand
      first = inst.calc
      6.times.map{inst.calc}.uniq.tap do |uniq_values|
        uniq_values.size.should == 1
        uniq_values.unify.should == first
      end
    end
  end

  describe '#clear_memos' do
    before :each do
      @method_names = %i[calc work apply]

      @inst = Class.new do
        include Memo
        memo def calc; rand; end
        memo def work; rand; end
        memo def apply; rand; end
      end.new

      method_names.each do |meth|
        @inst.send meth
      end
    end

    attr_reader :inst, :method_names

    it 'clears all memos' do
      method_names.each do |meth|
        inst.instance_variable_defined?("@#{meth}").should be_true
        inst.instance_variable_get("@#{meth}").should_not be_nil
      end

      inst.clear_memos

      method_names.each do |meth|
        inst.instance_variable_defined?("@#{meth}").should be_false
      end
    end

    it 'clears some memos' do
      method_names.each do |meth|
        inst.instance_variable_defined?("@#{meth}").should be_true
        inst.instance_variable_get("@#{meth}").should_not be_nil
      end

      to_clear = [:apply, :work]
      inst.clear_memos( *to_clear )

      to_clear.each do |meth|
        inst.instance_variable_defined?("@#{meth}").should be_false
      end

      inst.instance_variable_defined?("@calc").should be_true
    end
  end
end