Seq[(A, B)]の片側だけにflatMapするにはcats.arrowを使う

タプルのSeqの片側にflatMapしたいことがあって、arrowのsecondを使ったらうまくいったのでメモ。

flatMap(知ってる人は飛ばしていい)
タプルのSeqに対してflatMapしたい
- ナイーブな解法1: fをタプルに拡張する
- ナイーブな解法2: lisの片側にだけfを適用する
Arrow.secondを使った解法
- fをKleisliとして定義する
- cats.arrow.Arrowの拡張メソッドを使ってflatMap(f.second.run)する
全体のコード
まとめ
関連記事

flatMap(知ってる人は飛ばしていい)

皆さんflatMapって知ってますか？(YouTuber風) 例えば以下のようなことができて便利。

val lis = Seq(1, 2, 3, 4, 5)
val f: Int => Seq[Int] = n => Seq(n, n+1, n+2)
lis.flatMap(f) // => List(1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 5, 6, 5, 6, 7): scala.collection.immutable.Seq[scala.Int]

普通にmapするだけだと二重のSeqが作られてしまうのだが、flatMapすると1つのSeqに潰して平らにしてくれる。

flatMapする操作ができるのはSeqだけではなく、FutureやOptionといったF[_]といった形になるような多くの型でも使うことができる。より正確に関数型的な言い方で言うと、モナドのインスタンスになるようなF[_]ではflatMapを呼ぶことができる。

おおまかに言い換えると、f: A => F[B]になるような関数を受け取って、F[A]からF[B]に変換できるような操作が可能な型のことをモナドと言う*1。SeqもFutureもOptionもモナドなので、flatMapを呼び出せるようになっているというわけ。

タプルのSeqに対してflatMapしたい

さきほど自作プロダクトをいじっていて、こういう操作をしたいということがあった。

Seqに対してflatMapしたいのだが、flatMapしたい対象がペアに包まれてしまっている。つまりSeq[(A, B)]のBに対してflatMapがしたい。
渡す関数の型はB => Seq[C]、つまりペアの右側だけに作用する関数。
最終的にSeq[(A, C)]になっていてほしい。Cが増えた分は、単純にAを流用してほしい。

コードにするとこんな感じだ:

val lis: Seq[(String, Int)] = Seq(
  "a" -> 42,
  "b" -> 100,
  "c" -> 256,
)

val f: Int => Seq[Int] =  n => Seq(n, n+1, n+2)

// 最終的にSeq((a,42), (a,43), (a,44), (b,100), (b,101), (b,102), (c,256), (c,257), (c,258))が欲しい

ナイーブな解法1: fをタプルに拡張する

素朴にやると、f: (String, Int) => Seq[(String, Int)]になるようにfを書き直せばやれるのだが、fとは全く関係のない事情がfの実装に入り込むことになるので、メンテナンス困難になるし、なにより分かりにくい。

val f2: (String, Int) => Seq[(String, Int)] = {
  case (a, b) => Seq(a -> b, a -> b+1, a -> b+2)
}

ナイーブな解法2: lisの片側にだけfを適用する

lisをmapして、片側だけ剥がした状態でfを適用することもできるが、これだとどれだけ増えたのかが不明になってしまう。

val result = lis.map(_._2).flatMap(f)
lis.map(_._1) zip result // 個数が増えているので処理できない!

Arrow.secondを使った解法

このような操作は、CatsのArrowを使った関数の変形を使って簡単に実現できる。操作は以下のように行う:

fをKleisliとして定義する
cats.arrow.Arrowの拡張メソッドを使ってflatMap(f.second.run)する

これからそれぞれのステップについて解説する。

`f`を`Kleisli`として定義する

Catsが提供するArrowとは、矢印のような性質を持つ、合成が可能な概念をまとめてArrowという概念に統合し、合成や変換を行えるようにする仕組みである。例えばA => Bといった1引数関数もArrowとして合成や変換が可能である。

A => F[B]といった関数もArrowにすることができるのだが、型推論の関係でFの情報を型引数として取り出す必要があるので、一旦Kleisliというラッパーを挟む必要がある。Kleisliとは、A => F[B]の形になる関数のことで、要するにflatMapに入れるあの部分のことである。

といっても、書く上ではあまり難しいことはない。Kleisliの型引数は、F, 引数, 返り値の順に入れるだけだ。

import cats.data.Kleisli
val kf: Kleisli[Seq, Int, Int] = Kleisli { n => Seq(n, n+1, n+2) }

`cats.arrow.Arrow`の拡張メソッドを使って`flatMap(f.second.run)`する

CatsのArrowにはfirstとsecondという、A => Bの形になっているArrowを(A, X) => (B, X)や(X, A) => (X, B)の形に変形してくれるメソッドがある。これを使うことで、タプルの左側か右側かのどちらかにだけArrowを適用できるようになる。

typelevel.org

www.javadoc.io

そしてKleisliもArrowとして扱えるので、kf.secondと書くとKleisli[Seq, (X, Int), (X, Int)]の形になる。

最後に、runを呼び出すことでKleisil[Seq, (X, Int), (X, Int)]を元々の形である(X, Int) => Seq[(X, Int)]の形に戻し、flatMapに渡すことで完了する。

import cats.implicits._
lis.flatMap(kf.second.run) // => List((a,42), (a,43), (a,44), (b,100), (b,101), (b,102), (c,256), (c,257), (c,258))

全体のコード

import cats.implicits._
import cats.data.Kleisli

val lis: Seq[(String, Int)] = Seq(
  "a" -> 42,
  "b" -> 100,
  "c" -> 256,
)

val kf: Kleisli[Seq, Int, Int] = Kleisli { n => Seq(n, n+1, n+2) }

lis.flatMap(kf.second.run)