How can you create flexible enums with different behaviors?

Enums are great for grouping objects together that have similar behavior. They are also efficient because only one instance of the group will get created. However, it’s hard to implement enums for classes that have slightly different behaviors. Let me illustrate with an example.

Example: Stats calculator

Let’s say we want to write a stats calculator to compute mathematical statistics (e.g., SUM, COUNT, AVG, QUANTILES) for a list of values. Let’s first define an interface.

interface Stats{
  //Calcuate Stats for given list of values
  double calculate(List<Integer> values);
};

Now, we can easily represent different statistics as enum members like this:

enum Methods implements Stats{
  SUM {
      @Override
      public double calculate(List<Integer> values) {
          return (double) values.stream().mapToInt(i -> i).sum();
      }
  },
  COUNT {
      @Override
      public double calculate(List<Integer> values) {
          return (double) values.stream().count();
      }
  },
  AVG {
      @Override
      public double calculate(List<Integer> values) {
          return SUM.calculate(values) / COUNT.calculate(values);
      }
  }
}

We can use it like this:

public static void main(String[] args) {
  List<Integer> numbers = Arrays.asList(1,2,3,4,5,6,7,8,9,10);
  System.out.println(Methods.AVG.calculate(numbers));
  System.out.println(Methods.SUM.calculate(numbers));
}

Easy enough, but it becomes hard once we want to support the QUANTILE stat that needs to store the percentile to compute in an enum instance. For example, QUANTILE(90).calculate(values) should compute the 9090th percentile of values.

Problems with implementing enums for different behaviors

  • Since Java enums need to have the same fields, we have to add the percentile field to all enums even though it is irrelevant for AVG, SUM, and COUNT.

  • We don’t know the percentile values ahead of time to statically create instances of QUANTILE, so they can’t be enum anymore.

  • Once we convert them to regular classes, we will lose the default singleton behavior.

We basically want the simple statistics AVG, SUM, and COUNT as singletons, but QUANTILE as a dynamically created class for a given percentile value. We have to write a lot of boilerplate​ code to support this difference in behavior. Here is one way to do it:

public class StatsCalculator implements Stats {
  //To support singleton behaviour
  private static StatsCalculator SUM_INSTANCE = new StatsCalculator(SUM);
  private static StatsCalculator COUNT_INSTANCE = new StatsCalculator(COUNT);
  private static StatsCalculator AVG_INSTANCE = new StatsCalculator(AVG);
  private int percentile = 0;
  private Method method;
  @Override
  public double calculate(List<Integer> values) {
    switch (method){
      case AVG: return (double) values.stream().mapToInt(i -> i).sum() / values.stream().count()  ;
      case SUM: return (double) values.stream().mapToInt(i -> i).sum();
      case COUNT:return (double) values.stream().count();
      case QUANTILE: return Quantiles.percentiles().index(percentile).compute(values);
      default: throw new IllegalArgumentException("Invalid stat method");
    }
  }
  enum Method{
    SUM ,
    AVG,
    COUNT,
    QUANTILE;
  }
  private StatsCalculator(Method method, int percentile){
    this.method = method;
    this.percentile = percentile;
  }
  private StatsCalculator(Method method){
    this.method = method;
  }
  public static StatsCalculator getSimpleStats(Method method){
    switch (method){
        case AVG:return AVG_INSTANCE;
        case SUM:return SUM_INSTANCE;
        case COUNT:return COUNT_INSTANCE;
        default: throw new IllegalArgumentException("Should be called only for methods AVG OR SUM, OR COUNT");
    }
  }
  public static StatsCalculator getPercentileStats(Method method, int percentile){
    switch (method){
        case QUANTILE: return new StatsCalculator(method, percentile);
        default: throw new IllegalArgumentException("Should be called only for method QUANTILE");
    }
  }
  private double calculatePercentile(List<Integer> values){
    return Quantiles.percentiles().index(percentile).compute(values);
  }
  public static void main(String[] args) {
    List<Integer> numbers = Arrays.asList(1,2,3,4,5,6,7,8,9,10);
    System.out.println(getSimpleStats(AVG).calculate(numbers));
    System.out.println(getPercentileStats(QUANTILE,70).calculate(numbers));
  }
}

That’s lot of work.

Kotlin’s sealed classes to the rescue

Kotlin has powerful sealed classes to easily solve use cases like this.

Let’s first understand what a sealed class is:

Sealed classes are used to represent restricted class hierarchies. They are, in a sense, an extension of enum classes.

A sealed class can have subclasses, but all of them must be declared in the same file. A subclass of a sealed class can have multiple instances that can contain a state.

You can declare the subclasses inside or outside the sealed class, but they always have to be declared in the same file.

A sealed class is abstract by itself, it cannot be instantiated directly and can have abstract members. Sealed classes are not allowed to have non-private constructors (their constructors are private by default).

Now that we know about sealed classes, let’s see how we can use them to implement the Stats calculator.

Let’s create a sealed Stats class and subclasses for each Stats member, like this:

sealed class Stats {
  object Avg : Stats()    
  object Count : Stats()
  object Sum : Stats()
  class Quantile(val percentile: Int) : Stats() 
}

Kotlin supports object declaration to create singletons in a single line. The language itself supports defining exactly what we want:

  • The simple stats Avg, Count, and Sum are made singletons with object declaration.

  • Quantile is a regular class that stores a percentile value.

  • All of these classes are grouped together in the sealed base class Stats.

Now, the calculate method can be implemented in a clear and concise way:

fun calculate(values: List<Int>): Double {
  return calculate(this, values)
}
fun calculate(stats: Stats, values: List<Int>): Double =
  when (stats) {
    is Stats.Sum -> values.sum().toDouble()
    is Stats.Count -> values.size.toDouble()
    is Stats.Avg -> calculate(Stats.Sum, values) / calculate(Stats.Count, values)
    is Stats.Quantile -> Quantiles.percentiles().index(stats.percentile).compute(values)
}

The code is much more concise now. Furthermore, the beauty of the when expression helps us to avoid potential errors at runtime.

Since the when expression is used to directly return a value, if we add new stats, say Max extending Stats, but we forget to update the calculate method, then the compiler will throw an error.

The when expression is exhaustive, add the necessary ‘Max’ branch or ‘else’ branch instead.

This is possible with sealed classes because all subclasses are declared in the same file, so the compiler will know all possible values.

When are sealed classes useful?

The idea of sealed classes isn’t new. They allow us to easily work with algebraic data types. Similar features are available in other languages, like:

So, we can use Kotlin’s sealed classes to solve problems that require algebraic data types.

If you don’t have the luxury of using Kotlin, you can check out Spotify’s data enums in order to do the same work in Java.

Free Resources

License: Creative Commons-Attribution-ShareAlike 4.0 (CC-BY-SA 4.0)