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

-- iterateUnf :: Int -> a -> (a -> a) -> Unf a
-- iterateUnf n z f = Unf n z (\ _ x -> Just (x, f x))

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