Sun 01 November 2015

Power of Enums in Java

It comes sporadically that one has to implement a state machine for a project. Be it a benchmarking tool, or an authentication backend, or a Markov chain generator, it boils down to having a variable in one state, moving to another on certain action. In some of the cases, it might be quite useful to probabilisticaly move among states. There are many ways to implement such behaviour, but I found one really elegant, using Enums in Java.

The process outlined in this post will create a probabilistic, weighted state machine. This means that one could assign weights to states, and these weights will be used to determine the next state to switch to from the current state. I will rely on Enums, because they are elegant, they behave as classes, and the best of all, the system guarantees there won’t be multiple copies of the same state.

To make the process easier to follow, let’s assume we want to build a filesystem benchmark, where we issue different operations to a single file.

Since we deal with weighted probabilities, we first need an interface:

private interface WeightedChoice
{
    public int getWeight();
}

Then, we’ll implement the states as an implementation of WeightedChoice interface. We pass the weight in the constructor (unfortunately, one can’t define that via interfaces). To get the probabilities, the system will sum up all the weights given during the creation of the states.

private enum FileOperation implements WeightedChoice
{
    CREATE(10), // doesn't make much sense...
    EXISTS(5),
    READ(50),
    WRITE(20),
    GETATTR(10)
    DELETE(5);

    private int weight;

    private FileOperation(int weight)
    {
        this.weight = weight;
    }

    public int getWeight()
    {
        return weight;
    }
}

Et voila!

Now, the code that deals with the selection is a bit involved, but straightforward:

public static <T extends Enum<? extends T> & WeightedChoice> T choose(Class<T> clazz)
{
    return choose(ImmutableList.copyOf(clazz.getEnumConstants()));
}

public static <T extends Enum<? extends T> & WeightedChoice> T choose(List<T> from)
{
    final ThreadLocalRandom random = ThreadLocalRandom.current();
    final int[] stack = new int[from.size()];

    int weight = 0;
    int pos = 0;
    for (WeightedChoice elem : from)
    {
        weight += elem.getWeight();
        stack[pos++] = weight;
    }
    int x = random.nextInt(weight);
    for (int i = stack.length - 1; i >= 0; i--)
    {
        if (x > stack[i])
        {
            return from.get(i + 1);
        }
    }
    return from.get(0);
}

The above code will give us some state from the list (with the given probability distribution), but it can’t be used to create something like a Markov chain, where the next state depends on the current one. How to do that is a topic for another blog post.