{-# LANGUAGE ViewPatterns, TypeSynonymInstances, FunctionalDependencies, MultiParamTypeClasses, FlexibleInstances, UnboxedTuples, RankNTypes, BangPatterns, MagicHash, FlexibleContexts #-}

-- | A compilation of minor array combinators used extensively in "Data.RangeMin".
module Data.RangeMin.Internal.HandyArray (Producer, unsafeLookup, asPureArray, mListArray, mListUArray, mFuncArray, mFuncUArray, listMArray, listMUArray, listMArray', listMUArray', funcMArray, funcMUArray, MutArr(newArr, writeArr, readArr, lookup), MArr, MUArr, asMArr, asMUArr, mLook, muLook) where

import Data.RangeMin.Internal.HandyList
import Control.Monad
import Data.Array.Base
import Data.Array.IArray(Array)
import GHC.Arr(Ix(..))
import GHC.Prim
import GHC.Exts(Int(..))
import GHC.ST
import Control.Monad.State.Strict
import Control.Monad.Reader
import Data.Monoid
import Foreign.Storable(Storable(sizeOf))
import Prelude hiding (lookup)

type Producer m e = forall acc . (e -> acc -> m acc) -> acc -> m acc
type IA a s = StateT Int (Reader (a s)) (STBlank s)
data MArr e s = MArr Int# !(MutableArray# s e)
data MUArr s = MUArr Int# !(MutableByteArray# s)
type STBlank s = State# s -> State# s

boxIn :: (Int# -> e) -> Int -> e
boxIn f (I# i) = f i

blankIA :: IA a s
blankIA = return mempty

{-# INLINE fromBlank #-}
fromBlank :: STBlank s -> ST s ()
fromBlank m = ST $ \ s -> (# m s, () #)

instance Monoid (STBlank s) where
	{-# INLINE mempty #-}
	mempty s = s
	{-# INLINE mappend #-}
	mappend f g s = f (g s)

class MutArr a e | a -> e where
	newArr :: Int -> ST s (a s)
	readArr :: a s -> Int -> ST s e
	write :: a s -> Int -> e -> STBlank s
	writeArr :: a s -> Int -> e -> ST s ()
	lookup :: a s -> ST s (Int# -> e)

	{-# INLINE writeArr #-}
	writeArr arr !i x = fromBlank $ write arr i x

instance MutArr (MArr e) e where
	{-# INLINE newArr #-}
	newArr (I# n#) = ST $ \ s -> case newArray# n# (error "Undefined element") s of
		(# s', arr# #) -> (# s', MArr n# arr# #)
	{-# INLINE readArr #-}
	readArr (MArr n# arr#) (I# i#) = ST $ readArray# arr# i#
	{-# INLINE write #-}
	write (MArr n# arr#) (I# i#) = writeArray# arr# i#
	{-# INLINE lookup #-}
	lookup (MArr n# arr#) = ST $ \ s -> case unsafeFreezeArray# arr# s of
		(# s', iArr# #) -> (# s', \  i# -> case indexArray# iArr# i# of (# x #) -> x #)

instance MutArr MUArr Int where
	{-# INLINE newArr #-}
	newArr (I# n#) = ST $ \ s -> case newByteArray# (n# *# intSize#) s of
		(# s', arr# #) -> (# s', MUArr n# arr# #)
		where	intSize# = case sizeOf (1 :: Int) of I# i -> i
	{-# INLINE readArr #-}
	readArr (MUArr n# arr#) (I# i#) = ST $ \ s -> case readIntArray# arr# i# s of
		(# s', x# #) -> (# s', I# x# #)
	{-# INLINE write #-}
	write (MUArr n# arr#) (I# i#) (I# x#) = writeIntArray# arr# i# x#
	{-# INLINE lookup #-}
	lookup (MUArr n# arr#) = ST $ \ s -> case unsafeFreezeByteArray# arr# s of
		(# s', iArr# #) -> (# s', \ i# -> I# (indexIntArray# iArr# i#) #)

asPureArray :: Ix i => Array i e -> Array i e
asPureArray = id

{-# INLINE asMArr #-}
asMArr :: MArr e s -> MArr e s
asMArr = id

{-# INLINE asMUArr #-}
asMUArr :: MUArr s -> MUArr s
asMUArr = id

{-# INLINE decr #-}
decr :: (Monad m, MonadState Int m) => m Int
decr = modify (subtract 1) >> get

{-# INLINE extractor #-}
extractor :: ST s e -> (e -> STBlank s) -> STBlank s
extractor (ST k) f s = case k s of (# s', x #) -> f x s'

{-# SPECIALIZE INLINE acc :: e -> IA (MArr e) s -> IA (MArr e) s #-}
{-# SPECIALIZE INLINE acc :: Int -> IA MUArr s -> IA MUArr s #-}
acc :: MutArr a e => e -> IA a s -> IA a s
acc x ia = do	m <- ia
		!arr <- ask
		!j <- decr
		return (m `mappend` write arr j x)

{-# SPECIALIZE INLINE mAcc :: ST s e -> IA (MArr e) s -> IA (MArr e) s #-}
{-# SPECIALIZE INLINE mAcc :: ST s Int -> IA MUArr s -> IA MUArr s #-}
mAcc :: MutArr a e => ST s e -> IA a s -> IA a s
mAcc x ia = do	m <- ia
		!j <- decr
		!arr <- ask
		return (m `mappend` extractor x (write arr j))
--mapState (\ (m, i) -> let !j = i - 1 in  \ (ST x) (IA i m) -> let !j = i - 1 in IA j (\ s -> case x s of (# s', y #) -> m (write arr j y s'))

{-# SPECIALIZE INLINE accF :: MArr e s -> (Int -> e) -> Int -> STBlank s #-}
{-# SPECIALIZE INLINE accF :: MUArr s -> (Int -> Int) -> Int -> STBlank s #-}
accF :: MutArr a e => a s -> (Int -> e) -> Int -> STBlank s
accF arr f !n = foldr (mappend . acc') mempty (zeroEft (n-1)) where
	acc' !i = write arr i (f i)

{-# SPECIALIZE INLINE mAccF :: MArr e s -> (Int -> ST s e) -> Int -> STBlank s #-}
{-# SPECIALIZE INLINE mAccF :: MUArr s -> (Int -> ST s Int) -> Int -> STBlank s #-}
mAccF :: MutArr a e => a s -> (Int -> ST s e) -> Int -> STBlank s
mAccF arr f !n = foldr (mappend . mAcc') mempty (zeroEft (n-1)) where
	mAcc' !i = extractor (f i) (write arr i)

{-# INLINE unboxedLookup #-}
unboxedLookup :: IArray UArray e => UArray Int e -> (Int -> e)
unboxedLookup arr = unsafeAt $! arr

{-# INLINE mLook #-}
mLook :: MArr e s -> ST s (Int -> e)
mLook = liftM boxIn . lookup

{-# INLINE muLook #-}
muLook :: MUArr s -> ST s (Int -> Int)
muLook = liftM boxIn . lookup

runArr :: MutArr a e => Int -> (a s -> STBlank s) -> ST s (a s)
runArr !n f = do	!arr <- newArr n
			fromBlank $ f arr
			return arr

runIA :: MutArr a e => Int -> IA a s -> ST s (a s)
runIA !n m = runArr n $ runReader $ evalStateT m n

{-# INLINE mListArray #-}
mListArray :: Int -> [e] -> ST s (Int -> e)
mListArray n l = (runIA n $ foldr acc blankIA l) >>= mLook

{-# INLINE mListUArray #-}
mListUArray :: Int -> [Int] -> ST s (Int -> Int)
mListUArray n l = (runIA n $ foldr acc blankIA l) >>= muLook

{-# INLINE mFuncArray #-}
mFuncArray :: Int -> (Int -> e) -> ST s (Int -> e)
mFuncArray n f = mFuncArr n f >>= mLook

{-# INLINE mFuncUArray #-}
mFuncUArray :: Int -> (Int -> Int) -> ST s (Int -> Int)
mFuncUArray n f = mFuncArr n f >>= muLook

{-# INLINE mFuncArr #-}
mFuncArr :: MutArr a e => Int -> (Int -> e) -> ST s (a s)
mFuncArr !n f = runArr n (\ arr -> accF arr f n)

{-# INLINE listMArray #-}
listMArray :: Int -> [ST s e] -> ST s (Int -> e)
listMArray n l = (runIA n $ foldr mAcc blankIA l) >>= mLook

{-# INLINE listMUArray #-}
listMUArray :: Int -> [ST s Int] -> ST s (Int -> Int)
listMUArray n l = (runIA n $ foldr mAcc blankIA l) >>= muLook

{-# INLINE listMArray' #-}
listMArray' :: Int -> Producer (ST s) e -> ST s (Int -> e)
listMArray' n prod = listMArr' n prod >>= mLook

{-# INLINE listMUArray' #-}
listMUArray' :: Int -> Producer (ST s) Int -> ST s (Int -> Int)
listMUArray' n prod = listMArr' n prod >>= muLook

{-# INLINE listMArr' #-}
listMArr' :: MutArr a e => Int -> Producer (ST s) e -> ST s (a s)
listMArr' !n prod = do	arr <- newArr n
			prod (\ x i -> let !j = i - 1 in ST $ \ s -> (# write arr j x s, j #)) n
			return arr

{-# INLINE funcMArray #-}
funcMArray :: Int -> (Int -> ST s e) -> ST s (Int -> e)
funcMArray n f = funcMArr n f >>= mLook

{-# INLINE funcMUArray #-}
funcMUArray :: Int -> (Int -> ST s Int) -> ST s (Int -> Int)
funcMUArray n f = funcMArr n f >>= muLook

{-# INLINE funcMArr #-}
funcMArr :: MutArr a e => Int -> (Int -> ST s e) -> ST s (a s)
funcMArr !n f = runArr n (\ arr -> mAccF arr f n)

{-# INLINE unsafeLookup #-}
unsafeLookup :: (IArray a e, Ix i) => a i e -> i -> e
unsafeLookup arr@(bounds -> b) = unsafeAt arr . unsafeIndex b

{-# RULES
	"mconcat" forall l . mconcat l = foldr mappend mempty l;
	"$" forall f x . f $ x = f x;
-- 	"runST/liftM" forall (f :: e -> f) (m :: forall s . ST s e) . runST (liftM f m) = f (runST m)
	#-}