{-# LANGUAGE ViewPatterns, 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 Foreign.Storable(Storable(sizeOf))
import Prelude hiding (lookup)

type Producer m e = forall acc . (e -> acc -> m acc) -> acc -> m acc
data IA s = IA {-# UNPACK #-} !Int (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

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 acc #-}
acc :: MutArr a e => a s -> e -> IA s -> IA s
acc arr = \ x (IA i m) -> let !j = i - 1 in IA j (\ s -> m (write arr j x s))

{-# INLINE mAcc #-}
mAcc :: MutArr a e => a s -> ST s e -> IA s -> IA s
mAcc arr = \ (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 = acc' n where
	acc' !i	s	| i == 0	= s
			| otherwise 	= let !j = i - 1 in acc' j (write arr j (f j) s)

{-# INLINE mAccF #-}
mAccF :: MutArr a e => a s -> (Int -> ST s e) -> Int -> STBlank s
mAccF arr f !n = mAcc' n where
	mAcc' !i s	| i == 0	= s
			| otherwise	= let !j = i - 1 in case f j of 
						ST m -> case m s of 
							(# s', x #) -> mAcc' j (write arr j x s')

{-# 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
			case foldr (acc arr) (IA n initi) l of
				IA _ ans -> fromBlank ans >> mLook arr

{-# INLINE mListUArray #-}
mListUArray :: Int -> [Int] -> ST s (Int -> Int)
mListUArray n l = do	arr <- newArr n
			case foldr (acc arr) (IA n initi) l of 
				IA _ ans -> fromBlank ans >> 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
			case foldr (mAcc arr) (IA n initi) l of
				IA _ ans -> fromBlank ans >> mLook arr

{-# INLINE listMUArray #-}
listMUArray :: Int -> [ST s Int] -> ST s (Int -> Int)
listMUArray n l = do	arr <- newArr n
			case foldr (mAcc arr) (IA n initi) l of
				IA _ ans -> fromBlank ans >> 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