Kotlin の if else の評価値が直感と異なる件をコンパイラの視点で考察する

こんにちは、こんばんは、kitakkun です。

たまたま次のツイートを拝見しまして、面白いなと思ったので勝手ながら記事にまとめてみようと思います。

端的に述べますと、ここで問題として挙げられているコードは以下のようなものです。

if (condition1) { 
    ...
} else if (condition2) {
    ...
} else {
    ...
}.let { println(it) }

問題となっているコード

以下に、紹介されていたコードをそのまま引用します。

package syntax.weirdChaining
// by Kevin Most @kevinmost

fun printNumberSign(num: Int) {
  if (num < 0) {
    "negative"
  } else if (num > 0) {
    "positive"
  } else {
    "zero"
  }.let { println(it) }
}

printNumberSign(-2)
printNumberSign(0)
printNumberSign(2)

// What will it print?
// a) negative; zero; positive
// b) negative; zero
// c) negative; positive
// d) zero; positive

出典: https://github.com/angryziber/kotlin-puzzlers/blob/master/src/syntax/weirdChaining/WeirdChaining.kts

このプログラムは出力結果を予想するクイズとなっており、選択肢が4つ用意されています。

• a) negative; zero; positive
• b) negative; zero;
• c) nagative; positive
• d) zero; positive

おそらく多くの人は、引っ掛けだろうというのは分かりつつ、a) negative; zero; positive の順で出力されると考えるのではないでしょうか。私自身も正直最初は a かなと思いました。

以下のコードを Kotlin Playground にコピペして実行してみてください。きっと意外な結果が出るかと思います。

fun main() {
  printNumberSign(-2)
  printNumberSign(0)
  printNumberSign(2)
}

fun printNumberSign(num: Int) {
  if (num < 0) {
    "negative"
  } else if (num > 0) {
    "positive"
  } else {
    "zero"
  }.let { println(it) }
}

コードの実行結果

前節で紹介したコードを実行すると、次のような出力が得られるかと思います。

zero
positive

おや、negative が不在ですね。どういうことでしょうか。

コンパイラ視点からの考察

なぜ、negative の文字列のみ出力されなかったのでしょう。ここでは、コンパイラの内部処理に注目してその謎を読み解いていきたいと思います。なお、K2コンパイラ前提での考察となりますが、ご了承ください。

実は、Kotlin における if式 は、コンパイルの初期段階で when式 へとデシュガーされます。
実際に、Kotlin公式リポジトリ内に配置されている docs/fir/fir-basics.md にその記述があります。

• RAW_FIR: In this phase, we ran the translator from some parser AST to FIR tree. Currently, FIR supports two parser representations: PSI and LighterAst. During conversion, the FIR translator performs desugaring of the source code. This includes replacing all if expressions with corresponding when expressions, converting for loops into blocks with iterator variable declaration and while loop, and similar.

出典: Docs: add basic FIR compiler documentation · JetBrains/kotlin@d1335d3
※現在のバージョンでは該当の記述が削除されているため過去のコミット履歴を引用しています。

主にソースプログラムの解析を行うコンパイラFrontendにおいて、RAW_FIR を生成するタイミングで if を when に置き換えると説明されています。

このことから、おそらくですが、

if (num < 0) {
  "negative"
} else if (num > 0) {
  "positive"
} else {
  "zero"
}.let { println(it) }

という式は

when {
  num < 0 -> "negative"
  else -> when {
    num > 0 -> "positive"
    else -> "zero"
  }.let { println(it) }
}

のように内部的には扱われていることが予想されます。そして a) negative; zero; positive を選んだ我々が脳内でコンパイルした結果はこれとは別で、おそらく次のようなものを想像していたでしょう。

when {
  num < 0 -> "negative"
  num > 0 -> "positive"
  else -> "zero"
}.let { println(it) }

しかし、実際には .let の部分は全体にかかっているわけではなく、後半の if else の部分に対するものであったということになります。

仮説の検証

考えを表明することくらいは誰にでもできるので、実際に検証してみます。
幸い、手元に途中まで作成していた Kotlin Compiler Plugin の断片がありましたので、そちらを用いて実験を行いたいと思います。

問題となっているコードを、コンパイラプラグインを適用するプロジェクトの内部に配置します。

fun printNumberSign(num: Int) {
  if (num > 0) {
    "positive"
  } else if (num < 0) {
    "negative"
  } else {
    "zero"
  }.let { println(it) }
}

Firツリーが完全に解決されたタイミングでFirノードを確認したいので、FirAdditionalCheckersExtension を利用します。FirFunctionChecker を実装し、"printNumberSign" という名前の関数を見つけたら内容を吐き出すコードを書きます。

@OptIn(ExperimentalCompilerApi::class)
@AutoService(CompilerPluginRegistrar::class)
class MyCompilerRegistrar : CompilerPluginRegistrar() {
    override val supportsK2: Boolean get() = true

    override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
        FirExtensionRegistrarAdapter.registerExtension(MyFirExtensionRegistrar())
    }
}

@AutoService(AspectKFirExtensionRegistrar::class)
class MyFirExtensionRegistrar : FirExtensionRegistrar() {
    override fun ExtensionRegistrarContext.configurePlugin() {
        +::MyFirCheckerExtension
    }
}

class MyFirCheckerExtension(session: FirSession) : FirAdditionalCheckersExtension(session) {
    override val declarationCheckers = object : DeclarationCheckers() {
        override val functionCheckers: Set<FirFunctionChecker>
            get() = setOf(MyFunctionChecker())
    }
}

class MyFunctionChecker : FirFunctionChecker() {
    override fun check(declaration: FirFunction, context: CheckerContext, reporter: DiagnosticReporter) {
        if (declaration.symbol.name.asString() == "printNumberSign") {
            // めんどくさいのでエラーを吐いて表示させます
            error("${declaration.body?.render()}")
        }
    }
}

詳しいコンパイラプラグインの実装方法に関しては割愛しますが、最後にコンパイラプラグインのサンプルを紹介しますので気になる方はそちらを活用してご自身で試してみてください。

さて、作成したコンパイラプラグインを適用して該当の関数をコンパイルしてみます。error()で無理やりコンソールに吐いているので、エラーが出れば成功です。おそらくエラーメッセージの中に次のような文字列が見えるのではないでしょうか。

    when () {
        CMP(>, R|<local>/num|.R|kotlin/Int.compareTo|(Int(0))) ->  {
            String(positive)
        }
        else ->  {
            when () {
                CMP(<, R|<local>/num|.R|kotlin/Int.compareTo|(Int(0))) ->  {
                    String(negative)
                }
                else ->  {
                    String(zero)
                }
            }
            .R|kotlin/let|<R|kotlin/String|, R|kotlin/Unit|>(<L> = let@fun <anonymous>(it: R|kotlin/String|): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
                R|kotlin/io/println|(R|<local>/it|)
            }
            )
        }
    }

仮説でお見せしたコードと比較してみましょう。

when {
  num < 0 -> "negative"
  else -> when {
    num > 0 -> "positive"
    else -> "zero"
  }.let { println(it) }
}

確かに一致していますね！

まとめ

以上、Kotlin の if else の評価値が感覚とはちょっと違う不思議な例に関する考察でした。
問題になっているコードを書かれた Kevin Most さんは KotlinConf 2018 で 「Writing Your First Kotlin Compiler Plugin」というセッションをやられていた方です。さすが鋭い目をされています！

余談・宣伝

筆者はコンパイラを専門にやっていて、Kotlin Compiler Plugin を趣味で色々いじっております。大学の卒業研究では、コンパイラプラグインでデバッグ可能になるようにクラス内部を改変する back-in-time-plugin というものを制作しました。

また、現在は AspectJ 的なことを Kotlin Compiler Plugin を用いて実現する AspectK というプロジェクトを進めております。

もしご興味がありましたら、ご覧くださいませ！とてもシンプルな Kotlin Compiler Plugin の実装例も以下のリポジトリで公開しております！

ではでは！長々と失礼いたしました。皆様、良き Kotlin ライフを！