AtCoderでHaskell解がスタックオーバーフローを起こす問題の回避策

これを書いている時点でAtCoderHaskell処理系はGHC 7.4であり、これはプログラムに対してデフォルトで8MBのスタック制限を課す。したがって、深さ数十万の再帰を普通に書くと容易にスタック溢れによるREを起こす。

このスタック制限値はコンパイル時または実行時にRTSの-Kオプションで変更できるが、AtCoderではGHCにも自分のプログラムにも好きなオプションを渡すことが基本的にはできない*1ので、他の方法を考える。

回避策

以下のコードをコピペした上で、mainの先頭でunlimitStackSizeを呼べば良い。

import Control.Monad
import Data.Word
import Foreign.Ptr
import Foreign.Storable

-- | Remove the stack size limit in the RTS (+RTS -K).
unlimitStackSize :: IO ()
unlimitStackSize = do
  current <- peek maxStackSizePtr
  when (current == 0x100000) $ -- default limit: 1M words
    poke maxStackSizePtr (-1)

maxStackSizePtr :: Ptr Word32
maxStackSizePtr = plusPtr rtsFlagsPtr 12

foreign import ccall "&RtsFlags" rtsFlagsPtr :: Ptr Word32

RTSがオプションの値を保存するのに使っているグローバル変数の値を書き換えているだけだが、構造体内のメンバのオフセットなどをハードコードしているのでバージョンが変わると動作しなくなるかもしれない。もっとも、GHC 7.8でデフォルトのスタック制限が「物理メモリ量の80%」に変更されるので、おそらくこのようなハックは必要なくなるだろう。

(上記のコードは、この日記の他の部分と同様、コピー自由かつ無保証である。念のため)

*1:Template Haskellを使ってコンパイル時にGHCインスタンスをもう一つ起動するなどすれば可能かもしれない