{-# 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, listLookup, mListArray, mListUArray, mFuncArray, mFuncUArray, listMArray, listMUArray, listMArray', listMUArray', funcMArray, funcMUArray, MutArr(..), MArr, MUArr, asMArr, asMUArr, unboxedFreeze, mLook, muLook) where

import Data.RangeMin.Internal.HandyList
import Control.Monad
import Control.Monad.ST(ST, runST)
import Data.Array.ST(STArray, STUArray)
import Data.Array.IArray(Ix(..), IArray, Array)
import Data.Array.Base hiding (freeze)
import Data.Array.MArray hiding (freeze)
import Data.Array.Unboxed
import GHC.Arr(STArray(..), Ix(..))
import GHC.Exts
import GHC.Prim
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

instance Monoid (STBlank s) where
	mempty s = s
	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 = ST $ \ s -> (# write arr i x s, () #)

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

instance MutArr MUArr Int where
	{-# INLINE newArr #-}
	newArr (I# n#) {-| n# >=# 0#-} = 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#) {-| i# <# n# && i# >=# 0# -}= 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# -> {-if i# <# n# && i# >=# 0# then-} I# (indexIntArray# iArr# i#) {-else error "blahhh"-} #)

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 listLookup #-}
listLookup :: Int -> [e] -> (Int -> e)
listLookup n l = runST $ mListArray n l

{-# INLINE decr #-}
decr :: Monad m => StateT Int m e -> StateT Int m (Int, e)
decr = mapStateT (liftM (\ (x, i) -> let !j = i - 1 in ((j, x), j)))

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

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

{-# INLINE mAcc #-}
mAcc :: MutArr a e => ST s e -> IA a s -> IA a s
mAcc x ia = do	(j, m) <- decr ia
		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'))

{-# INLINE accF #-}
accF :: MutArr a e => a s -> (Int -> e) -> Int -> STBlank s
accF arr f n = mconcat $ map acc' $ zeroEft (n-1) where
	acc' = write arr `ap` f

{-# INLINE mAccF #-}
mAccF :: MutArr a e => a s -> (Int -> ST s e) -> Int -> STBlank s
mAccF arr f !n = mconcat $ map mAcc' $ zeroEft (n-1) where
	mAcc' = liftM2 extractor f (write arr)

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

{-# INLINE unboxedFreeze #-}
unboxedFreeze :: STUArray s Int Int -> ST s (UArray Int Int)
unboxedFreeze = unsafeFreeze

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

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

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

{-# INLINE mListArray #-}
mListArray :: Int -> [e] -> ST s (Int -> e)
mListArray n l = do	arr <- newArr n
			fromBlank $ runReader (evalStateT (foldr acc blankIA l) n) $! arr
			mLook arr

{-# INLINE mListUArray #-}
mListUArray :: Int -> [Int] -> ST s (Int -> Int)
mListUArray n l = do	arr <- newArr n
			fromBlank $ runReader (evalStateT (foldr acc blankIA l) n) $! arr
			muLook arr

{-# 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 = do	arr <- newArr n
			fromBlank $ accF arr f n
			return arr

{-# INLINE listMArray #-}
listMArray :: Int -> [ST s e] -> ST s (Int -> e)
listMArray n l = do	arr <- newArr n
			fromBlank $ runReader (evalStateT (foldr mAcc blankIA l) n) $! arr
			mLook arr

{-# INLINE listMUArray #-}
listMUArray :: Int -> [ST s Int] -> ST s (Int -> Int)
listMUArray n l = do	arr <- newArr n
			fromBlank $ runReader (evalStateT (foldr mAcc blankIA l) n) $! arr
			muLook arr

{-# 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 writeArr arr j x >> return 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 = do	arr <- newArr n
			fromBlank (mAccF arr f n)
			return arr

initi :: STBlank s
initi s = s

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