[Haskell-cafe] Lazy lists simulated by unboxed mutable arrays
Henning Thielemann
lemming at henning-thielemann.de
Fri May 30 14:59:19 EDT 2008
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)
More information about the Haskell-Cafe
mailing list