transform関数を使った効率化

 新しいStream Fusionでは,先に逐次処理を行う「map (+1) v // xs」のような処理に対しては,Stream FusionとRecyclingの両方の効率化を行えます。しかし,先に配列を更新する「map (+1) (v // xs)」のような処理に対しては,これだけでは十分に効率化できません。先に配列を更新する処理では,インライン化の結果として,"stream/unstream [Vector]"規則と"clone/new [Vector]"規則のどちらも適用できない式が生成されてしまうからです。

    map (+1) (v // xs)
==> { //演算子をインライン化 }
    map (+1)
        (new (New.modifyWithStream M.update
          (clone v) (Stream.fromList xs)))
==> { map関数のインライン化 }
    unstream (inplace (MStream.map (+1))
             (stream (new (New.modifyWithStream M.update
               (clone v) (Stream.fromList xs)))))
==> { unstream関数のインライン化 }
    (new
         (New.unstream (inplace (MStream.map (+1))
           (stream (new
               (New.modifyWithStream M.update
                   (clone v) (Stream.fromList xs)))))))

 最終的にインライン化された結果の式を見ると,問題はmap関数の「inplace (MStream.map f)」という処理を適用する部分にあることがわかります。「inplace (MStream.map f)」を適用する対象である「stream (new (New.modifyWithStream M.update ..))」に対しては,"stream/unstream [Vector]"規則も"clone/new [Vector]"規則も適用できません。このためStream Fusionを使った中間データ構造の除去も,Recyclingを使った配列の再利用もどちらも行われなくなります。

 map関数の定義を「New.unstream関数を適用し,その結果に対してmapを行う」ように変更すれば,この問題はとりあえず回避できます。しかし,それではmap関数に対してStream Fusionを行えなくなってしまいます。そこで,既存の書き換え規則を使うのではなく,「New.unstream (f (stream (new m)))」という式がどのような処理を行っているかに注目した新たな処理と書き換え規則を導入することにします。

 「New.unstream (f (stream (new m)))」という式は,New型の配列に「Stream a -> Stream b」型の関数fを適用してNew型の配列を返す処理になっています。したがってこの式は,new関数による配列の凍結やunstream関数による配列の作成を省いた別の処理に置き換えることができます。この処理を提供するのが,Data.Vector.Generic.Newモジュールのtransform関数です。

transform :: Vector v a =>
        (forall m. Monad m => MStream m a -> MStream m a) -> New v a -> New v a
{-# INLINE_STREAM transform #-}
transform f (New p) = New (MVector.transform f =<< p)

{-# RULES

"transform/transform [New]"
  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
         (g :: forall m. Monad m => MStream m a -> MStream m a)
         p .
  transform f (transform g p) = transform (f . g) p

"transform/unstream [New]"
  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
         s.
  transform f (unstream s) = unstream (f s)

 #-}

 transform関数は,New型の中身である可変配列を取り出し,Data.Vector.Generic.Mutableモジュールのtransform関数を使って関数fを可変配列に適用します。

-- Data.Vector.Generic.Mutable での transform 関数の定義
transform :: (PrimMonad m, MVector v a)
  => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
{-# INLINE_STREAM transform #-}
transform f v = fill v (f (mstream v))

mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
{-# INLINE mstream #-}
mstream v = v `seq` n `seq` (MStream.unfoldrM get 0 `MStream.sized` Exact n)
  where
    n = length v

    {-# INLINE_INNER get #-}
    get i | i < n     = do x <- unsafeRead v i
                           return $ Just (x, i+1)
          | otherwise = return $ Nothing

fill :: (PrimMonad m, MVector v a)
           => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
{-# INLINE fill #-}
fill v s = v `seq` do
                     n' <- MStream.foldM put 0 s
                     return $ unsafeSlice 0 n' v
  where
    {-# INLINE_INNER put #-}
    put i x = do
                INTERNAL_CHECK(checkIndex) "fill" i (length v)
                  $ unsafeWrite v i x
                return (i+1)

 Data.Vector.Generic.Mutableモジュールのtransform関数では,可変配列からStreamを作成するmstream関数と,Streamから可変配列を作成するfill関数に,同じNew型から取り出した可変配列を渡します。この結果,transform関数は「引数で受け取ったNew型の配列の中身を,関数fを使って直接書き換える処理」になるのです。このように,transform関数は,引数として取ったNew型の配列の中身を直接書き換えることで,new関数による配列の凍結やunstream関数による配列の作成を削除します。

 ただし,transform関数がNew型の配列の中身を直接書き換えるには,処理をすべてモナド上で行う必要があります。しかし,関数fのStream型に対する処理をモナド上で行うためには,transformを提供するだけでは不十分です。第40回で説明したように,Streamは「モナド上で処理を行う(=モナド化された)もの」にはなっていないからです。この問題を解決するために,vectorパッケージではモナド上で処理を行うためのMStream型を提供しています。

-- Data.Vector.Fusion.Stream での定義
-- | Alternative name for monadic streams
type MStream = M.Stream

-- Data.Vector.Fusion.Stream.Monadic での定義
-- | Result of taking a single step in a stream
data Step s a = Yield a s  -- ^ a new element and a new seed
              | Skip    s  -- ^ just a new seed
              | Done       -- ^ end of stream

-- | Monadic streams
data Stream m a = forall s. Stream (s -> m (Step s a)) s Size

 transform関数やmstream関数,fill関数の型宣言に現れていているMStream型は,「モナド化されたStream」だったのです。モナドを型変数として取らない従来のStreamは,MStreamを恒等モナドの上で利用するものとして,再定義されています。

-- Data.Vector.Fusion.Stream での定義
-- | The type of pure streams 
type Stream = M.Stream Id

-- | Identity monad
newtype Id a = Id { unId :: a }

instance Functor Id where
  fmap f (Id x) = Id (f x)

instance Monad Id where
  return     = Id
  Id x >>= f = f x

 以降では,「Data.Vector.Fusion.Stream.Monadicで提供されているStream型」と「Data.Vector.Fusion.Streamで提供されているStream型」を区別するため,「Data.Vector.Fusion.Stream.MonadicのStream」を「MStream」または「M.Stream」と表記します。単に「Stream」と表記した場合は,「Data.Vector.Fusion.StreamのStream(=MStream Id)」を指すことにします。

 vectorでは,第15回で説明したmtlパッケージのIdentity型ではなく,vectorのData.Vector.Fusion.Utilモジュールで独自に定義したId型の恒等モナドを使っています。mtlのIdentity型を使ったモナドとvectorのId型を使ったモナドは,型名やフィールド名こそ違うものの,恒等モナドであることに変わりはありません。

 mtlのIdentity型の恒等モナドではなく独自に定義したId型の恒等モナドを使っているのは,おそらくmtlへの依存を避けるためでしょう。vectorパッケージは,将来的にはGHCへの標準添付が想定されているため,すでにGHC標準添付から外されたmtlに依存することを避けたのだと考えられます(実際にGHCにvectorパッケージを含めるべきかどうかについては,2010年9月現在議論が行われています,参考リンク1参考リンク2)。

 話を戻しましょう。MStream型の導入により,transform関数を使ってStreamを利用する処理を再定義できるようになりました。しかし,それだけではまだ「New.unstream (f (stream (new m)))」を「transform f m」に書き換えることはできません。transform関数を使って書き換える前の「New.unstream (f (stream (new m)))」の関数fの型は「Stream a -> Stream b」であるのに対し,transform関数を使って書き換えた後の「transform f m」の関数fの型は「Monad m => MStream m a -> MStream m b」であり,それぞれの式で使われている関数fの型が異なるためです。

 第37回で説明したように,GHCの書き換え規則では変換元と変換先の型が等しい場合にのみ,式の書き換えを行います。このため,関数fの型が「MStream Id a -> MStream Id b」(すなわち「Stream a -> Stream b」)であるtrasform関数が恒等モナドに対して処理を行う場合にしか,transformを使った式への書き換えを行う書き換え規則が有効になりません。

 そこで「Monad m => MStream m a -> MStream m a」を「Stream a -> Stream b」に変換するinplace関数を使用し,fを「inplace f」として再定義します。

inplace :: (forall m. Monad m => M.Stream m a -> M.Stream m b)
        -> Stream a -> Stream b
{-# INLINE_STREAM inplace #-}
inplace f s = s `seq` f s

{-# RULES

"inplace/inplace [Vector]"
  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
         (g :: forall m. Monad m => MStream m a -> MStream m a)
         s.
  inplace f (inplace g s) = inplace (f . g) s

  #-}

 inplaceは関数fを使った多相的な処理を「Stream a -> Stream b」型に固定するだけのもので,他の処理は行いません。

 inplace関数の導入により,先の式は「New.unstream (inplace f (stream (new m)))」になります。「New.unstream (inplace f (stream (new m)))の関数fの型は「Monad m => MStream m a -> MStream m a」になり,「transform f m」に変換できるようになります。この変換を定義したものが,"inplace [Vector]"規則です。

{-# RULES
~ 略 ~

"inplace [Vector]"
  forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
  New.unstream (inplace f (stream (new m))) = New.transform f m

"uninplace [Vector]"
  forall (f :: forall m. Monad m => MStream m a -> MStream m a) m.
  stream (new (New.transform f m)) = inplace f (stream (new m))

 #-}

 "inplace [Vector]"規則を使ってtransform関数を利用する式に書き換えることで,「map (+1) (v // xs)」で行われる処理を,配列を1回しか複製しないよう効率化できます。

    map (+1) (v // xs)
==> { //演算子とmap関数の定義インライン化 }
    (new
         (New.unstream (inplace (MStream.map (+1))
           (stream (new
               (New.modifyWithStream M.update
                   (clone v) (Stream.fromList xs)))))))
==> { "inplace [Vector]"規則の適用 }
    new (New.transform (MStream.map (+1))
        (New.modifyWithStream M.update
            (clone v) (Stream.fromList xs)))

Id型の恒等モナドを使った関数の定義

 最後に,Id型の恒等モナドを使って定義されたStream型を使って処理を行う例として,vectorパッケージのfoldl関数やfoldr関数の定義を確認しておきましょう。

-- Data.Vector.Generic モジュールでの定義
-- | /O(n)/ Left fold
foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
{-# INLINE foldl #-}
foldl f z = Stream.foldl f z . stream

~ 略 ~

-- | /O(n)/ Right fold
foldr :: Vector v a => (a -> b -> b) -> b -> v a -> b
{-# INLINE foldr #-}
foldr f z = Stream.foldr f z . stream

-- Data.Vector.Fusion.Stream モジュールでの定義
-- | Left fold
foldl :: (a -> b -> a) -> a -> Stream b -> a
{-# INLINE foldl #-}
foldl f z = unId . M.foldl f z

~ 略 ~

-- | Right fold
foldr :: (a -> b -> b) -> b -> Stream a -> b
{-# INLINE foldr #-}
foldr f z = unId . M.foldr f z

 foldlやfoldrでは,Id型の恒等モナドを使って定義されたStream型を使って処理を行います。このStream型に対する処理は,Id型の恒等モナド上で行われるため,結果はId型に包まれたものになります。そこで,Data.Vector.Fusion.Streamのfoldlとfoldrでは,unIdフィールドを用いてStream型での計算結果である「Id b」型の値からb型の値を取り出すことで,最終的な結果を得ています。

 今回は,vectorパッケージで行われている書き換え規則を紹介しました。vectorパッケージでは,ほかにも効率化のためのさまざまな書き換え規則が提供されています。それらのほとんどは,特定の状況でしか使われないものです。こうした規則に興味がある方は,vectorパッケージのソースコードや,Recyclingについて説明した「Recycle Your Arrays!」という論文を参照してください。