[Haskell-cafe] Lazy lists simulated by unboxed mutable arrays

[Henning Thielemann]

On Wed, 28 May 2008, Ketil Malde wrote:

> Bulat Ziganshin <bulat.ziganshin at gmail.com> writes:
>
>> well, i don't understand difference between your idea and lazybs
>> implementation
>
> HT said earlier that:
>
>>> This would still allow the nice tricks for recursive Fibonacci
>>> sequence definition.
>
> Which I guess refers to something like:
>
>  fibs = 1 : 1 : zipWith (+) fibs (tail fibs)
>
> I don't think you can do that with LBS, since you'd need to calculate
> a whole chunk at a time, and for any chunk size > 1, each chunk
> depends on itself.

I have now implemented a small prototype:
   http://code.haskell.org/storablevector/Data/StorableVector/Cursor.hs

Actually you can run the Fibonacci example but it allocates three arrays:
   let f2 = zipNWith 15 (+) f0 f1; f1 = cons 1 f2; f0 = cons (0::Int) f1 in f0

I'm afraid the compiler cannot optimize the IORefs to unboxed values, even 
in registers, because in principle they can be modified from everywhere in 
the program. Is there a better way than using IORefs hidden by 
unsafePerformIO?



{-# OPTIONS_GHC -cpp -fglasgow-exts #-}
{- |
Simulate a list with strict elements by a more efficient array structure.
-}
module Data.StorableVector.Cursor where

import Control.Exception        (assert, )
import Control.Monad.State      (StateT(StateT), runStateT, )
import Data.IORef               (IORef, newIORef, readIORef, writeIORef, )

import Foreign.Storable         (Storable(peekElemOff, pokeElemOff))
import Foreign.ForeignPtr       (ForeignPtr, mallocForeignPtrArray, withForeignPtr, )

import Control.Monad            (when)
import Data.Maybe               (isNothing)

import System.IO.Unsafe         (unsafePerformIO)

import Data.StorableVector.Utility (viewListL, mapSnd, )

import Prelude hiding (length, foldr, zipWith, )


-- | Cf. StreamFusion  Data.Stream
data Generator a =
    forall s. -- Seq s =>
       Generator {
          next  :: {-# UNPACK #-} !(StateT s Maybe a),  -- compute next value
          state :: {-# UNPACK #-} !(IORef (Maybe s))    -- current state
       }

{- |
This simulates a
@ data StrictList a = Elem !a (StrictList a) | End @
by an array and some unsafe hacks.
-}
data Buffer a =
    Buffer {
        memory :: {-# UNPACK #-} !(ForeignPtr a),
        size   :: {-# UNPACK #-} !Int,  -- size of allocated memory
        gen    :: {-# UNPACK #-} !(Generator a),
        cursor :: {-# UNPACK #-} !(IORef Int)
    }

{- |
Vector is a part of a buffer.
-}
data Vector a =
    Vector {
        buffer :: {-# UNPACK #-} !(Buffer a),
        start  :: {-# UNPACK #-} !Int,   -- invariant: start <= cursor
        maxLen :: {-# UNPACK #-} !Int    -- invariant: start+maxLen <= size
    }


-- * construction

{-# INLINE create #-}
create :: (Storable a) => Int -> Generator a -> Buffer a
create l g = unsafePerformIO (createIO l g)

-- | Wrapper of mallocForeignPtrArray.
createIO :: (Storable a) => Int -> Generator a -> IO (Buffer a)
createIO l g = do
     fp <- mallocForeignPtrArray l
     cur <- newIORef 0
     return $! Buffer fp l g cur


{- |
@ unfoldrNTerm 20  (\n -> Just (n, succ n)) 'a' @
-}
unfoldrNTerm :: (Storable b) =>
    Int -> (a -> Maybe (b, a)) -> a -> Vector b
unfoldrNTerm i f x0 =
    unsafePerformIO (unfoldrNTermIO i f x0)

unfoldrNTermIO :: (Storable b) =>
    Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b)
unfoldrNTermIO i f x0 =
    do ref <- newIORef (Just x0)
       buf <- createIO i (Generator (StateT f) ref)
       return (Vector buf 0 i)


{-# INLINE pack #-}
pack :: (Storable a) => Int -> [a] -> Vector a
pack n = unfoldrNTerm n viewListL


{-# INLINE cons #-}
{- |
This is expensive and should not be used to construct lists iteratively!
-}
cons :: Storable a =>
    a -> Vector a -> Vector a
cons x xs =
    unfoldrNTerm (succ (maxLen xs))
       (\(mx0,xs0) ->
           fmap (mapSnd ((,) Nothing)) $
           maybe
              (viewL xs0)
              (\x0 -> Just (x0, xs0))
              mx0) $
    (Just x, xs)


{-# INLINE zipWith #-}
zipWith :: (Storable a, Storable b, Storable c) =>
    (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith f ps0 qs0 =
    zipNWith (min (maxLen ps0) (maxLen qs0)) f ps0 qs0


{-# INLINE zipNWith #-}
zipNWith :: (Storable a, Storable b, Storable c) =>
    Int -> (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipNWith n f ps0 qs0 =
    unfoldrNTerm n
       (\(ps,qs) ->
          do (ph,pt) <- viewL ps
             (qh,qt) <- viewL qs
             return (f ph qh, (pt,qt)))
       (ps0,qs0)



-- * inspection

-- | evaluate next value in a buffer
advanceIO :: Storable a =>
    Buffer a -> IO ()
advanceIO (Buffer p sz (Generator n s) cr) =
    do c <- readIORef cr
       assert (c < sz) $
          do writeIORef cr (succ c)
             ms <- readIORef s
             case ms of
                Nothing -> return ()
                Just s0 ->
                   case runStateT n s0 of
                      Nothing -> writeIORef s Nothing
                      Just (a,s1) ->
                         writeIORef s (Just s1) >>
                         withForeignPtr p (\q -> pokeElemOff q c a)

{-# INLINE switchL #-}
switchL :: Storable a => b -> (a -> Vector a -> b) -> Vector a -> b
switchL n j v = unsafePerformIO (switchLIO n j v)

switchLIO :: Storable a => b -> (a -> Vector a -> b) -> Vector a -> IO b
switchLIO n j v@(Vector buf st ml) =
    nullIO v >>= \ isNull ->
    if isNull
      then return n
      else
        do c <- readIORef (cursor buf)
           assert (st <= c) $ when (st == c) (advanceIO buf)
           x <- withForeignPtr (memory buf) (\p -> peekElemOff p st)
           let tl = assert (ml>0) $ Vector buf (succ st) (pred ml)
           return (j x tl)

{-# INLINE viewL #-}
viewL :: Storable a => Vector a -> Maybe (a, Vector a)
viewL = switchL Nothing (curry Just)


{-# INLINE foldr #-}
foldr :: (Storable a) => (a -> b -> b) -> b -> Vector a -> b
foldr k z =
    let recurse = switchL z (\h t -> k h (recurse t))
    in  recurse

{-# INLINE unpack #-}
unpack :: (Storable a) => Vector a -> [a]
unpack = foldr (:) []


instance (Show a, Storable a) => Show (Vector a) where
    showsPrec p x = showsPrec p (unpack x)


{-# INLINE null #-}
null :: Vector a -> Bool
null = unsafePerformIO . nullIO

nullIO :: Vector a -> IO Bool
nullIO (Vector (Buffer _ sz (Generator _ s) _) st _) =
    do b <- readIORef s
       return (st >= sz || isNothing b)