「関数型デザイン＆プログラミング」でのScalaでのST sモナドの実装について

書籍「関数型デザイン&プログラミング」をマイペースに読み進めています。
この書籍の14章では、局所作用を実行するためのクラスであるSTとSTRefが紹介されています。
今回はこれについてまとめてみます。

ST、STRefは以下のImmutableとみなされるための条件を満たし、状態遷移を局所化する仕組みを提供するクラスです。状態遷移をローカルスコープに閉じたものとし、それをScalaの型システムを利用して強制させることで、参照透過性を保証させます。

Immutableな型とみなすための条件

以下の条件に違反するものはコンパイルされないようにScalaの型システムを利用します。

• ミュータブルなオブジェクトへの参照を保持しているもの以外は、そのオブジェクトが変更されていることを認識できない。
• ミュータブルなオブジェクトを、それが作成されたスコープの外側から参照することができない。

（「関数型デザイン&プログラミング」より引用）

STの実装

上記を踏まえたSTの実装は以下のようになっています。
(Github: fpinscala/answers/src/main/scala/fpinscala/localeffects/LocalEffects.scala
より)

sealed trait ST[S,A] { self => 
    protected def run(s: S): (A,S) //protectedにすることにより、runが実行可能なスコープを制限。変更を外部に漏らさないようにする

    def map[B](f: A => B): ST[S,B] = new ST[S,B] { 
        def run(s: S) = { 
            val (a, s1) = self.run(s) 
           (f(a), s1) 
        } 
    } 
    def flatMap[B](f: A => ST[S,B]): ST[S,B] = new ST[S,B] { 
        def run(s: S) = { 
            val (a, s1) = self.run(s) 
            f(a).run(s1) 
        } 
    } 
} 

object ST { 
    def apply[S,A](a: => A) = { 
        lazy val memo = a //runが複数回呼ばれることを考慮したキャッシュ
        new ST[S,A] { 
            def run(s: S) = (memo, s) 
        } 
    } 

    def runST[A](st: RunnableST[A]): A = 
        st[Unit].run(())._1 
} 

上記で注目したいのはSTのrunメソッドです。runメソッドをprotectedで修飾することにより、外部からのアクセスを制限しています。これがpublicだと、引数のsを外部から渡せるようになってしまいます。runメソッド内で使われる可能性のあるSを外部で保持していることになるので、上記の不変条件に違反してしまいます。
また、runメソッドのシグネチャは S => (A, S)なので、runによりA型の値が新たに生成されることを示しています。

※runSTについては後述。

STRefの実装

STRefはST内の計算で使用されるミュータブルな参照を表現しています。STRefをST内のみでしか参照・操作できないよう実装することにより、状態遷移を局所化します。上記の不変条件を満たすよう実装されているため、ST内でSTRefがもつ値を変更したとしても、それが外部に漏れることはありません。
上記を踏まえたSTの実装は以下のようになっています。
(Github: fpinscala/answers/src/main/scala/fpinscala/localeffects/LocalEffects.scala
より)

sealed trait STRef[S,A] { //sealedで修飾することにより、外部からSTRefを生成することはできない
    protected var cell: A 
    def read: ST[S,A] = ST(cell) 
    def write(a: => A): ST[S,Unit] = new ST[S,Unit] { 
        def run(s: S) = { 
        cell = a 
        ((), s) 
        } 
    } 
} 

object STRef { 
    def apply[S,A](a: A): ST[S, STRef[S,A]] = ST(new STRef[S,A] { 
        var cell = a 
    }) 
} 

STRefはsealedにより修飾されているので、外部から直接STRefを生成することはできません。

sealed trait STRef[S,A] { //sealedで修飾することにより、外部からSTRefを生成することはできない

外部からSTRefを生成するには、STRef.apply呼び出す必要がありますが、これはSTによりラッピングされた形で返却されます。

  //applyメソッド呼び出しにより、STRefはSTにラッピングされた形で生成される
  //ST[S, A]とSTRef[S, A]のSの型を一致させる
  def apply[S,A](a: A): ST[S, STRef[S,A]] = ST(new STRef[S,A] { 

ここで注目したい点は以下です。

・applyで生成されるST[S, A]とSTRef[S, A]のSの型が一致している
これにより、型Sを利用して、ST内でSTRefが操作できることになります。

・STRefを外部から直接生成できず、STの内部にラップされる形でしか生成できないようにしている
これにより、STRefの生存期間はSTの生存期間に束縛されます。
STRefの状態遷移をSTの中に閉じ込めることで、外部からは副作用のないようにものとして振る舞うことが可能となります。

・apply、read、writeメソッドは戻り値としてSTを返す

 //以下の３つのメソッドはSTを戻り値として返すため、STの計算が継続できる
 def apply[S,A](a: A): ST[S, STRef[S,A]]
 def read: ST[S,A]
 def write(a: => A): ST[S,Unit]

上記のようにapply、read、writeは戻り値としてSTを返します。
これにより、上記演算はSTのコンテキストに紐づくようになるため、STRefを操作する計算が継続できるようになります。

ST[S, A]とSTRef[S, A]の型パラメータSについて

お気づきかもしれませんが、上記STとSTRefの実装では型Sを使用していません。
このSはSTRefをST内で操作するためのトークンのような役割しか持っていません。
なので、このSの型は実のところ何でも良いということになります。
これは、STのアクションが型Sに関して多相であると考えられます。

RunnableSTの定義

上記のSTの実装ではrunSTメソッドが定義されていました。その実装を再掲すると以下のようになっています。

object ST {
  （中略）
  def runST[A](st: RunnableST[A]): A = 
        /*
        val as = st[Unit].run(()) //as: (A, S)
        val a = as._1 // a: A
        */
        st[Unit].run(())._1 //Sは多相であるため、Unitを指定しても問題ない。
}

引数にRunnableSTなるものが指定されています。
STアクションからSTRefを取り出せないように、Scalaの型システムを利用して定義したものがRunnableSTです。

RunnableSTの実装は以下のようになっています。

(Github: fpinscala/answers/src/main/scala/fpinscala/localeffects/LocalEffects.scala
より)

trait RunnableST[A] {
  def apply[S] : ST[S, A]
}

RunnableSTを作成する目的は以下です。

1. ST[S, STRef[S, A]]型のアクションが実行できないようにする
外部からST[S, STRef[S, A]]のアクションの実行を許可しまうと、その結果としてSTRef[S, A]が取得できることになってしまいます。状態遷移をローカルスコープ内に局所化するためSTRefが外部にさらされることになるので、これを防ぎます。

2. T型がS型に関与するようなST[S, T] のアクションが実行できないようにする
S型の変更により、T型も変更されるようなアクション、
T型の変更により、S型も変更されるようなアクションは
同様に状態遷移の局所化に反するので、これを実行できないよう防ぎます。

3. ST[S, A]の型Sが多相であることを表現する
上述の通り、ST・STRefでの型Sはトークンの役割しか持っていません。
この型Sが多相であることを表現します。

STのrunメソッドはprotectedで宣言されているため、外部からSTをrunするためには、このrunSTメソッドを使用しなければなりません。

具体的な使用例は以下のような感じです。

val p = new RunnableST[Int] { //A型だけを指定し、S型を指定する必要がない
  /*
  型Sがapplyにバインドしているため、new RunnableSTを指定した場合、
  型Sにアクセスすることはできない
  => 型SにSTRefを指定して、STRefを外部から取得することはできない
  */
  def apply[S] = for { 
    r <- STRef(1)
    x <- r.read
    _ <- r.write(x + 1)
    y <- r.read
  } yield y
}

val r = ST.runST(p) // r: Int = 2

このように、STRefをSTアクションの内側だけにとどめ、参照を変化させても安全な仕組みを提供しています。綺麗ですね。