{-# LANGUAGE ExistentialQuantification, BangPatterns #-}
module Data.RangeMin.Common.Unf (generateUnf, postscanlUnf', toUnf, unfold, unfoldM, unfoldInto, Unf(..)) where

import Control.Monad
-- import Control.Monad.ST
import Control.Monad.Primitive
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import Data.RangeMin.Common.Vector

data Unf a = forall b . Unf {-# UNPACK #-} !Int b (Int -> b -> Maybe (a, b))

generateUnf :: Int -> (Int -> a) -> Unf a
generateUnf n f = Unf n () (\ i _ -> Just (f i, ()))

postscanlUnf' :: (a -> b -> a) -> a -> Unf b -> Unf a
postscanlUnf' f z0 (Unf n s0 suc) = 
	Unf n (z0, s0) (\ i (z, s) -> do
		(x, s') <- suc i s
		let !z' = f z x
		return (z', (z', s')))

{-# INLINE toUnf #-}
toUnf :: G.Vector v a => v a -> Unf a
toUnf !xs = generateUnf (G.length xs) (xs !)

{-# INLINE unfold #-}
unfold :: G.Vector v a => Unf a -> v a
unfold unf = inlineCreate (unfoldM unf)

{-# INLINE unfoldM #-}
unfoldM :: (GM.MVector v a, PrimMonad m) => Unf a -> m (v (PrimState m) a)
unfoldM unf@(Unf n _ _) = do
	!dest <- new n
	unfoldInto0 dest unf
	return dest

-- {-# INLINE unfoldSnocM #-}
-- unfoldSnocM :: (GM.MVector v a, PrimMonad m) => Unf a -> a -> m (v (PrimState m) a)
-- unfoldSnocM unf@(Unf n _ _) z = do
-- 	!dest <- new (n+1)
-- 	unfoldIntoSnoc0 dest unf z
-- 	return dest

{-# INLINE unfoldInto0 #-}
unfoldInto0 :: (GM.MVector v a, PrimMonad m) => v (PrimState m) a -> Unf a -> m ()
unfoldInto0 !dest (Unf n s0 suc) = do
	let go i s = when (i < n) $ case suc i s of
		Nothing -> return ()
		Just (x, s') -> write dest i x >> go (i+1) s'
	go 0 s0
	return ()

-- {-# INLINE unfoldIntoSnoc0 #-}
-- unfoldIntoSnoc0 :: (GM.MVector v a, PrimMonad m) => v (PrimState m) a -> Unf a -> a -> m ()
-- unfoldIntoSnoc0 !dest (Unf n s0 suc) z = do
-- 	let go i s = if (i < n) then case suc i s of
-- 		Nothing -> write dest i z
-- 		Just (x, s') -> write dest i x >> go (i+1) s'
-- 	      else write dest i z
-- 	go 0 s0
-- 	return ()

{-# INLINE unfoldInto #-}
unfoldInto :: (G.Vector v a, PrimMonad m) => G.Mutable v (PrimState m) a -> Unf a -> m (v a)
unfoldInto dest unf@(Unf n _ _) = do
	unfoldInto0 dest unf
	unsafeFreeze (sliceM 0 n dest)

-- {-# INLINE unfoldIntoSnoc #-}
-- unfoldIntoSnoc :: (G.Vector v a, PrimMonad m) => G.Mutable v (PrimState m) a -> Unf a -> a -> m (v a)
-- unfoldIntoSnoc dest unf@(Unf n _ _) z = do
-- 	unfoldIntoSnoc0 dest unf z
-- 	unsafeFreeze (sliceM 0 n dest)