ListとStreamの両方でheadとtailにmatchさせる

Listを扱うプログラムを書く時に、match式で::を使ってheadとtailを取り出すことが良くあると思う。
例えば、「二つのList a, bがあり、bがaに前方一致するならaからbを取り除いた残りを戻す」というような処理をmatchを使って書くと*1、こんな風になるかと思う。

   def dropWith[A](a: List[A], b: List[A]): Option[List[A]] = {
      (a, b) match {
         case (_, Nil) => Some(a)
         case (Nil, _) => None
         case (h1 :: _, h2 :: _) if h1 != h2 => None
         case (h1 :: t1, h2 :: t2) => dropWith(t1, t2)
      }
   }

scala> dropWith(List(1, 2, 3, 4, 5), List(1, 2))
res1: Option[List[Int]] = Some(List(3, 4, 5))

scala> dropWith(List(1, 2, 3, 4, 5), List(1, 3))
res2: Option[List[Int]] = None

これを拡張して、一つ目の引数が無限リスト(というかStream)でもOKにしたい。それで一つ目の引数の型をSeqにしてみたんだけど、List以外を指定すると「::」でマッチさせようとしてもうまく動かない。「::」に加えて「#::」も書いておくとStreamの場合にも対応できるけど、表現が冗長だし、結局ListとStream以外(例えばRange)が来た場合に動かない。*2

   // scalaVersion := "2.9.1"
   def dropWith[A](a: Seq[A], b: List[A]): Option[Seq[A]] = {
      (a, b) match {
         case (_, Nil) => Some(a)
         case (Nil, _) => None
         case (h1 :: _, h2 :: _) if h1 != h2 => None
         case (h1 :: t1, h2 :: t2) => dropWith(t1, t2)
         case (h1 #:: _, h2 :: _) if h1 != h2 => None
         case (h1 #:: t1, h2 :: t2) => dropWith(t1, t2)
      }
   }

scala> dropWith(Stream.from(1), List(1, 2))
res1: Option[Seq[Int]] = Some(Stream(3, ?))

scala> dropWith(1 to 5, List(1, 2))
scala.MatchError: (Range(1, 2, 3, 4, 5),List(1, 2)) (of class scala.Tuple2)
...

どうしたらいいかなーと思って検索してると、extractorを自分で定義している例を見つけた。

scala> val fff = new { def unapply[A, T](x : { def head : A; def tail : T;
def isEmpty : Boolean }) = if(x.isEmpty) None else Some(x.head, x.tail) }
fff: java.lang.Object{def unapply[A,T](AnyRef{def head: A; def tail: T;
def isEmpty: Boolean}): Option[(A, T)]} = $anon$1@32b90a0

scala> fff.unapply(List(1,2))
res112: Option[(Int, List[Int])] = Some((1,List(2)))

scala> List(1,2,3) match { case x fff xs => (x, xs

)}

res113: (Int, List[Int]) = (1,List(2, 3))

scala> Stream(1,2,3) match { case x fff xs => (x, xs)}
res114: (Int, Stream[Int]) = (1,Stream(2, ?))

Using the foldLeft operator /: | The Scala Programming Language

上の例では構造的部分型を使っているけど、ドキュメントを見るとtailは共通の親traitであるSeqLikeに定義されているので、SeqLikeに対してextractorを用意してやれば良いと思う。

   import collection.SeqLike
   def dropWith[A, S <% SeqLike[A, S]](a:S, b:Seq[A]): Option[S] = {
      object ::#:: {
         def unapply[T <% SeqLike[A, T]](x: T) = if (x.isEmpty) None else Some(x.head, x.tail) 
      }  
      (a, b) match {
         case (_, Nil) => Some(a)
         case (Nil, _) => None
         case (h1 ::#:: _, h2 ::#:: _) if h1 != h2 => None
         case (h1 ::#:: t1, h2 ::#:: t2) => dropWith(t1, t2)
      }  
   }

scala> dropWith(List(1, 2, 3, 4, 5), List(1, 2))
res0: Option[List[Int]] = Some(List(3, 4, 5))

scala> dropWith(Stream.from(1), List(1, 2))
res1: Option[scala.collection.immutable.Stream[Int]] = Some(Stream(3, ?))

scala> dropWith("pineapple", "pine")
res2: Option[java.lang.String] = Some(apple)

scala> dropWith(1 to 5, 1 to 3)
<console>:11: error: No implicit view available from scala.collection.immutable.Range.Inclusive => scala.collection.SeqLike[Int,scala.collection.immutable.Range.Inclusive].
              dropWith(1 to 5, 1 to 3)
                      ^

うーん、Rangeでは結局上手くいかない。RnageはSeqLike[Int, Range]ではなくSeqLike[Int, IndexedSeq[Int] ]を継承しているせいみたい。

ちなみにStreamが::をサポートしていない理由は「Scala by Example」に書いてあって、::がtailをListとして事前に評価してしまうから(遅延評価を旨とするStreamとは合わない)らしい。

追記: やっぱり構造的部分型で受けてみた

   object :#: {
      def unapply[A, T <% {def head: A; def tail: T; def isEmpty: Boolean}](x: T): Option[(A, T)] =
          if (x.isEmpty) None else Some(x.head, x.tail)
   }  
   @tailrec
   def dropWith[A, T <% {def head: A; def tail: T; def isEmpty: Boolean}](a: T, b: Seq[A]): Option[T] = {
      (a, b) match {
         case (_, Nil) => Some(a)
         case (Nil, _) => None
         case (h1 :#: _, h2 :#: _) if h1 != h2 => None
         case (h1 :#: t1, h2 :#: t2) => dropWith(t1, t2)
      }
   }

scala> dropWith(Stream.from(1), List(1, 2))
res0: Option[scala.collection.immutable.Stream[Int]] = Some(Stream(3, ?))

scala> dropWith("pineapple", "pine")
res1: Option[java.lang.String] = Some(apple)

scala> dropWith(1 to 5, 1 to 3)
<console>:11: error: No implicit view available from scala.collection.immutable.Range.Inclusive => AnyRef{def head: Int; def tail: scala.collection.immutable.Range.Inclusive; def isEmpty: Boolean}.
              dropWith(1 to 5, 1 to 3)
                      ^

scala> dropWith[Int,Range](1 to 5, 1 to 3)
res3: Option[Range] = Some(Range(4, 5))

*1:お話の都合上matchを使っているので、if (a startsWith b) Some(a drop b.length) else Noneで良いじゃんというのは無しで……

*2:まああと戻り値が共変じゃないので使いにくい。