#ifndef MEMOIZE_H_INCLUDED
#define MEMOIZE_H_INCLUDED

#include <map>
#include <utility>

/* Utility classes for defining memoized recursive
 * functions. To define a unary recursive function
 * rec:
 *
 *   class Rec : public UnaryMemoFunction< Arg, Result >
 *   {
 *      Result call( Arg a ) { ... }
 *   } rec;
 *
 * To define a binary recursive function rec:
 *
 *   class Rec : public BinaryMemoFunction< Arg1, Arg2, Result >
 *   {
 *      Result call( Arg1 a, Arg2 b ) { ... }
 *   } rec;
 *
 * In both cases, call() should look like the normal
 * recursive code except it should call memo(...) to
 * call itself.
 *
 * Use size() to see how many calls have been memoized.
 * Use clear() to erase memoized calls.
 */

template < class Arg, class Result >
class UnaryMemoFunction
{
private:
  typedef std::map< Arg, Result > Cache;

public:
  Result operator() ( Arg a ) { return memo( a ); }

  void clear() { cache.clear(); }
  size_t size() { return cache.size(); }

protected:
  Result memo( Arg a )
  {
    typename Cache::const_iterator it = cache.find( a );
    return it == cache.end() ? cache[ a ] = call( a ): it->second;
  }

  virtual Result call( Arg n ) = 0;

private:
  Cache cache;
};


template < class Arg1, class Arg2, class Result >
class BinaryMemoFunction
{
private:
  typedef std::pair< Arg1, Arg2 > Pair;
  typedef std::map< Pair, Result > Cache;

public:
  Result operator() ( Arg1 a, Arg2 b ) { return memo( a, b ); }

  void clear() { cache.clear(); }
  size_t size() { return cache.size(); }

protected:
  Result memo( Arg1 a, Arg2 b )
  {
    Pair p( a, b );
    typename Cache::const_iterator it = cache.find( p );
    return it == cache.end() ? cache[ p ] = call( a, b ): it->second;
  }

  virtual Result call( Arg1 a, Arg2 b ) = 0;

private:
  Cache cache;
};
#endif // MEMOIZE_H_INCLUDED