Stream.hs

{-# LANGUAGE CPP          #-}
{-# LANGUAGE BangPatterns #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Data.Array.Vector.Strict.Stream
-- Copyright   : (c) 2006 Roman Leshchinskiy
-- License     : see libraries/ndp/LICENSE
-- 
-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
-- Stability   : internal
-- Portability : non-portable (existentials)
--
-- Description ---------------------------------------------------------------
--
-- Stream combinators and fusion rules for flat unboxed arrays.
--


#include "fusion-phases.h"

module Data.Array.Vector.Strict.Stream (
  streamU, unstreamU, unstreamMU
) where

import Data.Array.Vector.Prim.Hyperstrict (
    (:*:)(..), fstS, sndS)
import GHC.ST (ST)
import Data.Array.Vector.Stream (
  Step(..), Stream(..), mapS, zipWithS)
import Data.Array.Vector.UArr (
  UArr, MUArr, UA, indexU, lengthU, zipU, fstU, sndU, newDynU, writeMU)

-- | Generate a stream from an array, from left to right
--
streamU :: UA a => UArr a -> Stream a
{-# INLINE_STREAM streamU #-}
streamU !arr = Stream next 0 n
  where
    n = lengthU arr
    {-# INLINE next #-}
    next i | i == n    = Done
           | otherwise = Yield (arr `indexU` i) (i+1)

-- | Create an array from a stream, filling it from left to right
--
unstreamU :: UA a => Stream a -> UArr a
{-# INLINE_STREAM unstreamU #-}
unstreamU st@(Stream next s n) = newDynU n (\marr -> unstreamMU marr st)

-- | Fill a mutable array from a stream from left to right and yield
-- the number of elements written.
--
unstreamMU :: UA a => MUArr a s -> Stream a -> ST s Int
{-# INLINE_U unstreamMU #-}
unstreamMU marr (Stream next s n) = fill s 0
  where
    fill s i = i `seq`
               case next s of
                 Done       -> return i
                 Skip s'    -> s' `seq` fill s' i
                 Yield x s' -> s' `seq`
                               do
                                 writeMU marr i x
                                 fill s' (i+1)




-- | Fusion rules
-- --------------

-- The main fusion rule

{-# RULES

"streamU/unstreamU" forall s.
  streamU (unstreamU s) = s

  #-}

-- Zip fusion
--
-- NB: We do not separate rules for zip3U etc. because these are implemented
-- in terms of zipU

{-# RULES

"streamU/zipU" forall a1 a2.
  streamU (zipU a1 a2) = zipWithS (:*:) (streamU a1) (streamU a2)

"fstU/unstreamU" forall s.
  fstU (unstreamU s) = unstreamU (mapS fstS s)
"sndU/unstreamU" forall s.
  sndU (unstreamU s) = unstreamU (mapS sndS s)

  #-}