{-# OPTIONS_GHC -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, withForeignPtr, )
-- import Foreign.Ptr              (Ptr)
import Data.StorableVector.Memory (mallocForeignPtrArray, )

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

import System.IO.Unsafe         (unsafePerformIO, unsafeInterleaveIO, )

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

import Prelude hiding (length, foldr, zipWith, )


-- | Cf. StreamFusion  Data.Stream
data Generator a =
   forall s. -- Seq s =>
      Generator
         {-# UNPACK #-} !(StateT s Maybe a)  -- compute next value
         {-# 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 l f x0 =
   unsafePerformIO (unfoldrNTermIO l f x0)

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

unfoldrN :: (Storable b) =>
   Int -> (a -> Maybe (b, a)) -> a -> (Vector b, Maybe a)
unfoldrN l f x0 =
   unsafePerformIO (unfoldrNIO l f x0)

unfoldrNIO :: (Storable b) =>
   Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b, Maybe a)
unfoldrNIO l f x0 =
   do ref <- newIORef (Just x0)
      buf <- createIO l (Generator (StateT f) ref)
      s <- unsafeInterleaveIO $
             do evaluateToIO l buf
                readIORef ref
      return (Vector buf 0 l, s)
{-
unfoldrNIO :: (Storable b) =>
   Int -> (a -> Maybe (b, a)) -> a -> IO (Vector b, Maybe a)
unfoldrNIO l f x0 =
   do y <- unfoldrNTermIO l f x0
--      evaluateTo l y
      let (Generator _ ref) = gen (buffer y)
      s <- readIORef ref
      return (y, s)

Data/StorableVector/Cursor.hs:98:10:
    My brain just exploded.
    I can't handle pattern bindings for existentially-quantified constructors.
    In the binding group
        (Generator _ ref) = gen (buffer y)
    In the definition of `unfoldrNIO':
        unfoldrNIO l f x0
                     = do
                         y <- unfoldrNTermIO l f x0
                         let (Generator _ ref) = gen (buffer y)
                         s <- readIORef ref
                         return (y, s)
-}


{-
unfoldrN :: (Storable b) =>
   Int -> (a -> Maybe (b, a)) -> a -> (Vector b, Maybe a)
unfoldrN i f x0 =
   let y = unfoldrNTerm i f x0
   in  (y, getFinalState y)

getFinalState :: (Storable b) =>
   Vector b -> Maybe a
getFinalState y =
   unsafePerformIO $
      ...
-}


{-# 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

-- zipWith f ps qs = pack $ List.zipWith f (unpack ps) (unpack qs)

{-# 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)
{-
let f2 = zipNWith 15 (+) f0 f1; f1 = cons 1 f2; f0 = cons (0::Int) f1 in f0

*Data.StorableVector.Cursor> let xs = unfoldrNTerm 20  (\n -> Just (n, succ n)) (0::Int)
*Data.StorableVector.Cursor> let ys = unfoldrNTerm 20  (\n -> Just (n, 2*n)) (1::Int)
*Data.StorableVector.Cursor> zipWith (+) xs ys
-}




-- * 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)

-- | evaluate all values up to a given position
evaluateToIO :: Storable a =>
   Int -> Buffer a -> IO ()
evaluateToIO l buf@(Buffer _p _sz _g cr) =
   whileM
      (fmap (<l) (readIORef cr))
      (advanceIO buf)

whileM :: Monad m => m Bool -> m a -> m ()
whileM p f =
   let recurse =
          do b <- p
             when b (f >> recurse)
   in  recurse

{-# 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

-- | /O(n)/ Converts a 'Vector a' to a '[a]'.
{-# 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)
--   assert (l >= 0) $ l <= 0


{-
toVector :: Storable a => Vector a -> VS.Vector a
toVector v =
   VS.Cons (memory (buffer v)) ()
-}

-- length