Kotlin when(string) optimization

TL;DR

• Kotlin’s when(string) semantics allow flexibility, but default to if/else bytecode implementation when non-constants are involved, causing extraneous string equality checks
• You can return the optimization by separating constant and non-constant strings to different whens
• This optimization is a net benefit and thus should be in the compiler

Background

Unciv is pretty hyper-optimized by now, where hashmap lookups and String comparisons — even with cached hashes — feel far too slow.

I recently encountered that compiler optimization in Java 7 generates hashcode-based switching from String switch-cases, and recalled we have a function with a big String-based switch that’s a hotspot.

Given a 30s CPU profiling run, we see this function running ~500ms of a 30s run — in other words, 1.6% runtime — checking string comparisons. This is very noticeable for the first options, where we return a static constant value.

The Bytecode

Let’s take a look at just that first line, "Terrain" -> true , in IDEA’s Kotlin Bytecode tool (double-shift — ‘show Kotlin bytecode`) — annotated for readability.

L0
  LINENUMBER 492 L0
  ALOAD 1 // store the filter..
  ASTORE 3 // in register 3
 L1
  LINENUMBER 493 L1
  ALOAD 3 // load register 3 to the stack
  LDC "Terrain" // push the string to compare to the stack, and compare
  INVOKESTATIC kotlin/jvm/internal/Intrinsics.areEqual (Ljava/lang/Object;Ljava/lang/Object;)Z
  IFEQ L2 // if value is 0 (false) go to next case
  ICONST_1 // correct case! Load 1 (true) 
  GOTO L3 // go to 'return' 
 L2
  LINENUMBER 494 L2
 FRAME APPEND [java/lang/String]
  ALOAD 3
  LDC "All"
  INVOKESTATIC kotlin/jvm/internal/Intrinsics.areEqual (Ljava/lang/Object;Ljava/lang/Object;)Z
  IFEQ L4
  ICONST_1
  GOTO L5
 L4
 FRAME SAME
  ALOAD 3
  LDC "all"
  INVOKESTATIC kotlin/jvm/internal/Intrinsics.areEqual (Ljava/lang/Object;Ljava/lang/Object;)Z
 L5
 FRAME SAME1 I
  IFEQ L6
  ICONST_1
  GOTO L3
 L6
...

We can see we’re sequentially checking each option using String equality — not great.

What if we take the first couple of cases and put them in another switch?

return when (filter) {
    "Terrain" -> true
    "All", "all" -> true
    else -> false
}

Turns into:

LINENUMBER 492 L0
ALOAD 1 // 
ASTORE 3 // store filter in register 3
ALOAD 3
INVOKEVIRTUAL java/lang/String.hashCode ()I // get hashcode
LOOKUPSWITCH // switch by hashcode! optimized! :D
  65921: L1 // in compile time we generate the hashcode for each option
  96673: L2
  241216277: L3
  default: L4
L1 // "All" hashcode case
FRAME APPEND [java/lang/String]
ALOAD 3 // Push filter to stack 
LDC "All" // Push string to compare, and compare (below)
INVOKEVIRTUAL java/lang/String.equals (Ljava/lang/Object;)Z 
IFNE L5 // if true, goto 'return true'
GOTO L4 // else goto 'return false'
L2
FRAME SAME
ALOAD 3
LDC "all"
INVOKEVIRTUAL java/lang/String.equals (Ljava/lang/Object;)Z
IFNE L5
GOTO L4
L3
FRAME SAME
ALOAD 3
LDC "Terrain"
INVOKEVIRTUAL java/lang/String.equals (Ljava/lang/Object;)Z
IFEQ L4 // if false goto 'return false'
L6
LINENUMBER 493 L6
ICONST_1
GOTO L7
L5 // 'return true' label
LINENUMBER 494 L5
FRAME SAME
ICONST_1 // push 'true' to stack
GOTO L7 // goto 'return'
L4 // 'return false' label
LINENUMBER 495 L4
FRAME SAME
ICONST_0
L7 // return value in stack
LINENUMBER 492 L7
FRAME SAME1 I
IRETURN
L8

Here we see the optimization in play, we initially find the hash of the string and then use a JVM lookupswitch on the hash vs precomputed hashes of the options to direct us to the only possible correct option — so we only run String.equals() once!

What makes it break is when we introduce a compiler unknown comparison, like

return when (filter) {
    "Terrain" -> true
    "All", "all" -> true
    runtimeVariable -> false
}

The compiler cannot determine the value of all hashcodes at runtime, so it defaults to if/else equality. :(

Attempts to theoretically optimize can also harm you. The naive thought is, “if we have several strings that lead to the same result, instead of adding many cases let’s have a stored set of them and check if the set contains the value” — but instead, we ruin the compiler’s built in optimization for the entire when block!

Solution

We can solve this by splitting the original when into 2: one for constants, and the other for non-constants.

       return when(string){
            "a" -> 1
            "b" -> 2
            "c" -> 3
            runtimeVariable -> 4
            "d" -> 5
            else -> 6
        }

becomes

      return when(string){
            "a" -> 1
            "b" -> 2
            "c" -> 3
            else -> when(string){
                runtimeVariable -> 4
                "d" -> 5
                else -> 6
            }
        }

The Result

As a consequence of splitting it up, we went from ~500ms / 30s to ~40ms / 30s — in other words we shaved off 90% of the work, and saved 1.5% of runtime!

This belongs in the compiler

This problem is unique to Kotlin, since in Java it won’t let you use non-constant values in switch statements. The amazing versatility that when gives you works against you in edge cases that don’t fit Java’s notion of what a switch looks like. But since this is a guaranteed performance improvement, this should be part of the compilers job: separate cases into constant and non-constant, turn the constants into a lookup and the non-constants into if/elses.

And so, it became an official compiler suggestion :)