Java memoization method

Question

I came across an interesting problem and was wondering if and how could this be done in Java: Create a method which can memoize any function/method . The method has the following arguments : the method/function and the argument(s) for it.

For example let's say i have this method :

int addOne(int a) { return a + 1;}

and i call my memoization method two times with the same arguments : addOne and 5 for example, the first call should actually call the addOne method and return the result and also store that result for that given argument. The second time when i call it should know this has been called before and just look up the previous answer.

My idea would be to have something like a HashMap<Callable,HashMap<List<Objects>,Object>> where you would store the previous answers and look them up later on.I think this can be somehow done with lambda expressions but i'm not that familiar with them.I'm not quite sure how to write this method and would appreciate some help.

Can this be done with this approach?

Answer 1

In Java 8 you can use ConcurrentHashMap.computeIfAbsent:

Map<Integer, Integer> cache = new ConcurrentHashMap<>();

Integer addOne(Integer x) {
    return cache.computeIfAbsent(x -> x + 1);
}

DZone has a good tutorial which provides a solution that will work for any method:

The Memoizer class is quite simple:

public class Memoizer<T, U> {

  private final Map<T, U> cache = new ConcurrentHashMap<>();

  private Memoizer() {}

  private Function<T, U> doMemoize(final Function<T, U> function) {
    return input -> cache.computeIfAbsent(input, function::apply);
  }

  public static <T, U> Function<T, U> memoize(final Function<T, U> function) {
    return new Memoizer<T, U>().doMemoize(function);
  }
}

Using this class is also extremely simple:

Integer longCalculation(Integer x) {
  try {
    Thread.sleep(1_000);
  } catch (InterruptedException ignored) {
  }
  return x * 2;
}
Function<Integer, Integer> f = this::longCalculation;
Function<Integer, Integer> g = Memoizer.memoize(f);

public void automaticMemoizationExample() {
  long startTime = System.currentTimeMillis();
  Integer result1 = g.apply(1);
  long time1 = System.currentTimeMillis() - startTime;
  startTime = System.currentTimeMillis();
  Integer result2 = g.apply(1);
  long time2 = System.currentTimeMillis() - startTime;
  System.out.println(result1);
  System.out.println(result2);
  System.out.println(time1);
  System.out.println(time2);
}

Running the automaticMemoizationExample method will produce the following result:

2
2
1000
0

Answer 2

You can memoize any function with Java 8's MethodHandles and lambdas if you're willing to give up type safety on the parameters:

public interface MemoizedFunction<V> {
    V call(Object... args);
}

private static class ArgList {
    public Object[] args;

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (!(o instanceof ArgList)) {
            return false;
        }

        ArgList argList = (ArgList) o;

        // Probably incorrect - comparing Object[] arrays with Arrays.equals
        return Arrays.equals(args, argList.args);
    }

    @Override
    public int hashCode() {
        return args != null ? Arrays.hashCode(args) : 0;
    }
}

public static <V> MemoizedFunction<V> memoizeFunction(Class<? super V> returnType, Method method) throws
                                                                                                  IllegalAccessException {
    final Map<ArgList, V> memoizedCalls = new HashMap<>();
    MethodHandles.Lookup lookup = MethodHandles.lookup();
    MethodHandle methodHandle = lookup.unreflect(method)
                                      .asSpreader(Object[].class, method.getParameterCount());
    return args -> {
        ArgList argList = new ArgList();
        argList.args = args;
        return memoizedCalls.computeIfAbsent(argList, argList2 -> {
            try {
                //noinspection unchecked
                return (V) methodHandle.invoke(args);
            } catch (Throwable throwable) {
                throw new RuntimeException(throwable);
            }
        });
    };
}

Working Example

This creates a variable-arity lambda that encloses the function and is almost as fast as calling the function directly (i.e., no reflection happens inside of call(Object...args)) after the lambda is constructed since we're using MethodHandle.invoke() instead of Method.invoke().

You can still do this without lambdas (replace with anonymous classes) and MethodHandles (replace with Method.invoke), but there will be performance penalties that make this less attractive for performance-conscious code.