「次に何を実行するかの判定」と「次を実行する責任」を分離して、StackOverflowErrorに立ち向かう

ヌーラボでScalaを書くRubyistの谷本です。
先日、初めて Play Framework に送ったプルリクエストがマージされました🎉

このプルリクエストは該当ロジックでStackOverflowErrorが起こらないように修正するものです。今回はどのように考えて修正したか説明しようと思います。

StackOverflowErrorとは

StackOverflowErrorとは名前の通りStackがあふれたため起こるエラーです。
詳細の説明のため、Stackとは何かを見てみましょう。今回は、ScalaコードだったのでJVMのStackについてドキュメントを参照します。OracleのドキュメントのJava Virtual Machine Stacksという項目にはStackはframeを保持するものであると書かれています。さらにframeはメソッドの返り値などを保持するためのもので、メソッドを呼び出した時に作られ、メソッド完了時に破棄されると書かれています。(JVMにはNative Method Stacksもありますが今回は割愛します。)
例えば、再起呼び出しのようにメソッドの中でメソッドを呼び出し、さらにそのメソッドの中でメソッドを呼び出すというコードを書くと、その呼び出し階層分だけStackにframeが積まれます。処理対象データによっては呼び出し階層が深くなりすぎ、JVMに設定されたStackサイズでは耐えられなくなると、StackOverflowErrorが起こります。
このため、StackOverflowErrorを回避するには呼び出し階層が深くなりすぎないように設計するのが重要です。

プルリクエストで解決したかった問題

実はこのプルリクエストは2年前に自分が報告したStackOverflowErrorについてのIssueを修正するもので、当時は設定で回避できると言われてそのままにしていました。ただ、これを設定で回避しようとするとなかなか大変なので、そもそもStackOverflowErrorが起こらないようにしたのが今回のプルリクエストです。
対象コードはパーサーの実装で、あるメソッドmethodAからmethodBかmethodCを、それぞれのメソッドからmethodAや他のメソッドを呼び出すといった具合でいくつかのメソッドが組み合わった再帰呼び出しになっています。
簡単に書くと以下のようなコードだと思ってください。

def methodA(argument: Argument): Unit = {
  someLogicA // main logic of methodA

  // execute the next method.
  if (someConditionA) {
    methodB(argument)
  } else {
    methodC(argument)
  }
}

def methodB(argument: Argument): Unit = {
  someLogicB // main logic of methodB

  // execute the next method.
  if (someConditionB) {
    methodC(argument)
  } else {
    methodA(argument)
  }
}

def methodC(argument: Argument): Unit = {
  someLogicC // main logic of methodC

  // execute the next method.
  if (someConditionC) {
    methodA(argument)
  }
}

def main(argument: Argument): Unit = {
  methodA(argument)
}

修正方法

再帰呼び出しのStackOverflowErrorを回避する常套手段としては、ループに書き換えるというものでしょう。Scalaの場合は末尾再帰にしてしておくとコンパイラ側でループに書き換えてくれるので、末尾再帰に書き換えるという手もあります。よくあるパターンとしては以下のような書き換えです。(以下の例はすでに末尾再帰になっているのでScalaでは修正の必要がありませんが。)

再帰呼び出しでの実装

def mainFunc(n: Int): Unit = {
  if (n > 0) {
    // Execute some logic in methodA
    yourImportantLogic1(n)
    yourImportantLogic2(n)

    methodA(n - 1)
  } else {}
}

ループでの実装

def mainFunc(n: Int): Unit = {
  // Execute some logic in methodA
  yourImportantLogic1(n)
  yourImportantLogic2(n)
}

def loopMethod(n: Int): Unit = {
  var i = n
  while(i > 0) {
    methodA(i)
    i = i - 1
  }
}

この書き換えは「次に何を実行するかの判定」と「次を実行する責任」をループ側に移動することととらえることができると思います。上の例では、「n > 0の時n - 1を使ってmainFuncを呼び出さなければならない」が「次に何を実行するかの判定」で、「実際にmethodAを呼び出す」が「次を実行する責任」です。
ただ、「次に何を実行するかの判定」は元のメソッドであるmainFuncに残して置いて、「次を実行する責任」だけをループ側に移動したいときもあります。例えば、パーサーなどは実行中のメソッドから次のメソッドが決まるようなパターンが多いと思うので、コード上はyourImportantLogic1/2と「次に何を実行するかの判定」は同じメソッド内で呼び出す方が読みやすいことがあります。
これを実現する一つの方法としては、「次に何を実行するかの判定」をした結果をlambda式で返し、ループの方でそのlambda式の実行を行うという方法です。こうすることで「次に何を実行するかの判定」はmethodAに残すことができ、元のコードの意図をあまり変えない書き換えができます。
これを「プルリクエストで解決したかった問題」の節に書いた簡単なコードに適用してみるとどうなるか見てみましょう。

trait Result
case class Continue(next: () => Result) extends Result
object Finished extends Result

def methodA(argument: Argument): Result = {
  someLogicA // main logic of methodA

  // Make information to execute the next method.
  if (someConditionA) {
    Continue(() => methodB(argument))
  } else {
    Continue(() => methodC(argument))
  }
}

def methodB(argument: Argument): Result = {
  someLogicB // some logic of methodB

  // Make information to execute the next method.
  if (someConditionB) {
    Continue(() => methodC(argument))
  } else {
    Continue(() => methodA(argument))
  }
}

def methodC(argument: Argument): Result = {
  someLogicC // some logic of methodB

  // Make information to execute the next method.
  if (someConditionC) {
    Continue(() => methodA(argument))
  } else Finished
}

def main(argument: Argument): Unit = {
  var result = methodA(argument)
  while(result != Finished) {
    result = result.asInstanceOf[Continue].next()
  }
}

if (someConditionA)のような「次に何を実行するかの判定」がmethodA/methodB/methodCの各メソッドにとどまり、各メソッドはContinueクラスでラップしたlambda式を返すか完了を表すFinishedを返すようになりました。そのlambda式はmainメソッド内のwhileループで実行されるようになり「次を実行する責任」が移動しています。それぞれのメソッドはContinueかFinishedを返して完了するのでStack内のframeも破棄されStackOverflowErrorが起こらないようになっています。
また、mainメソッド内を以下のように書き換えることで末尾再帰のコードにすることもできます。

def main(argument: Argument): Unit = {
  @tailrec
  def loop(result: Result): Unit = {
    result match {
      case Continue(next) => loop(next())
      case Finished => ()
    }
  }
  loop(methodA(argument))
}

終わりに

今回は、lambda式を利用して、「次に何を実行するかの判定」と「次を実行する責任」を分離してStackOverflowErrorを回避する方法を紹介しました。複雑そうなコードでも意外に簡単に対応できるので一つの方法として覚えておくと良いのではないかと思います。また、一見すると分割できなさそうにも見える「次に何を実行するかの判定」と「次を実行する責任」を分割するという考え方を持っていればプログラミングの幅も広がるかもしれません。
最後になりましたが、Lightbend社のCLA(Contributor Licence Agreement)のチェックをしてくれたNulab OSPOの皆さんありがとうございました。この場を借りてお礼申し上げます。(現在はCLAに同意する必要はなくなったようです。)
では、よいOSS生活を！