    {- |
    Module      :  Control.Monad.Trans.Memo.Array
    Copyright   :  (c) Eduard Sergeev 2013
    License     :  BSD-style (see the file LICENSE)

    Maintainer  :  eduard.sergeev@gmail.com
    Stability   :  experimental
    Portability :  non-portable (multi-param classes, flexible instances)

    ArrayCache - mutable-array-based (`IO` and `ST` hosted) `MonadCache`

    Very fast memoization cache. Unfortunatelly it cannot suit every case (see limitations), but if you can use it, please do: it is generally an order of magnitude faster than `Data.Map`-based `Control.Monad.Trans.Memo.Map.Memo`, especially /unboxed/ version - try to use it whenever you can.

    Limitations: Since `Data.Array.Base.MArray` is used as `MonadCache` the key range must be known beforehand and the array is allocated before the first call.
    It is therefore most suitable for the cases when the distribution of possible key values is within reasonable range and is rather dense (the best case: all values withing some range will be used). If this is the case then `MArray` has O(1) for both lookup and update operations.
    In addition unboxed `UArrayCache` can only store unboxed types (but it does it very efficiently).

    -}

    {-# LANGUAGE NoImplicitPrelude,
      MultiParamTypeClasses, FunctionalDependencies,
      FlexibleInstances, FlexibleContexts, TypeSynonymInstances,
      UndecidableInstances, TypeFamilies #-}

    module Control.Monad.Memo.Array
    (

       -- * ArrayCache for boxed types
       Array,
       ArrayCache,
       ArrayMemo,
       evalArrayMemo,
       runArrayMemo,
       -- * ArrayCache for unboxed types
       UArray,
       UArrayCache,
       UArrayMemo,
       evalUArrayMemo,
       runUArrayMemo,
       -- * Generic function for ArrayCache
       Container(..),
       Cache,
       genericEvalArrayMemo,
       genericRunArrayMemo

    ) where

    import Data.Function
    import Data.Maybe (Maybe(..))
    import Data.Array.ST
    import Data.Array.IO
    import Control.Monad
    import Control.Monad.Trans.Class
    import Control.Monad.ST
    import System.IO

    import Data.MaybeLike
    import Control.Monad.Memo.Class
    import Control.Monad.Trans.Memo.ReaderCache

    newtype Container arr = Container { forall arr. Container arr -> arr
    toArray :: arr }

    -- | Generic Array-based memo cache
    type Cache arr k e = ReaderCache (Container (arr k e))

    instance (Monad m, Ix k, MaybeLike e v, MArray c e m) =>
        MonadCache k v (Cache c k e m) where
            {-# INLINE lookup #-}
            lookup :: k -> Cache c k e m (Maybe v)
    lookup k
    k = do
              c <- ReaderCache (Container (c k e)) m (Container (c k e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              e <- lift $ readArray (toArray c) k
              return (if isNothing e then Nothing else Just (fromJust e))
            {-# INLINE add #-}
            add :: k -> v -> Cache c k e m ()
    add k
    k v
    v = do
              c <- ReaderCache (Container (c k e)) m (Container (c k e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              lift $ writeArray (toArray c) k (just v)

    instance (Monad m, Ix k, MaybeLike e v, MArray c e m) =>
        MonadMemo k v (Cache c k e m) where
            {-# INLINE memo #-}
            memo :: (k -> Cache c k e m v) -> k -> Cache c k e m v
    memo k -> Cache c k e m v
    f k
    k = do
              c <- ReaderCache (Container (c k e)) m (Container (c k e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              e <- lift $ readArray (toArray c) k
              if isNothing e
                then do
                  v <- f k
                  lift $ writeArray (toArray c) k (just v)
                  return v
                else return (fromJust e)

    -- ArrayCache for boxed types
    -- --------------------------

    -- | A family of boxed arrays
    type family Array (m :: * -> *) :: * -> * -> *

    type instance Array (ST s) = STArray s
    type instance Array IO = IOArray

    type instance Array (ReaderCache c (ST s)) = STArray s
    type instance Array (ReaderCache c IO) = IOArray

    -- | Memoization monad based on mutable boxed array
    type ArrayCache k e m = Cache (Array m) k e m

    -- | This is just to be able to infer the type of the `ArrayCache` element
    --
    -- Type families could be used instead but due to the bug in 7.4.* we cannot use them here
    class MaybeLike e v => ArrayMemo v e | v -> e

    -- | Evaluate computation using boxed array
    --
    -- Key range should cover all possible keys used in computation
    -- otherwise /not in range/ error is generated by array
    evalArrayMemo :: (Ix k, MArray (Array m) e m, ArrayMemo v e) =>
                     ArrayCache k e m a  -- ^memoized computation to be evaluated
                  -> (k,k)  -- ^array key range
                  -> m a  -- ^computation result
    {-# INLINE evalArrayMemo #-}
    evalArrayMemo :: forall k (m :: * -> *) e v a.
    (Ix k, MArray (Array m) e m, ArrayMemo v e) =>
    ArrayCache k e m a -> (k, k) -> m a
    evalArrayMemo = Cache (Array m) k e m a -> (k, k) -> m a
    forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m a
    genericEvalArrayMemo

    -- | Evaluate computation and the final content of array cache using boxed array
    --
    -- Key range should cover all possible keys used in computation
    -- otherwise /not in range/ error is generated by array
    runArrayMemo :: (Ix k, MArray (Array m) e m, ArrayMemo v e) =>
                    ArrayCache k e m a  -- ^memoized computation to be evaluated
                 -> (k,k)  -- ^array key range
                 -> m (a, Array m k e)  -- ^computation result and final array cache
    {-# INLINE runArrayMemo #-}
    runArrayMemo :: forall k (m :: * -> *) e v a.
    (Ix k, MArray (Array m) e m, ArrayMemo v e) =>
    ArrayCache k e m a -> (k, k) -> m (a, Array m k e)
    runArrayMemo = Cache (Array m) k e m a -> (k, k) -> m (a, Array m k e)
    forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m (a, arr k e)
    genericRunArrayMemo

    -- ArrayCache for unboxed types
    -- ----------------------------

    -- | A family of unboxed arrays
    type family UArray (m :: * -> *) :: * -> * -> *

    type instance UArray (ST s) = STUArray s
    type instance UArray IO = IOUArray

    type instance UArray (ReaderCache c (ST s)) = STUArray s
    type instance UArray (ReaderCache c IO) = IOUArray

    -- | Memoization monad based on mutable unboxed array
    type UArrayCache k e m = Cache (UArray m) k e m

    -- | This is just to be able to infer the type of the `UArrayCache` element
    --
    -- Type families could be used instead but due to the bug in 7.4.* we cannot use them here
    class MaybeLike e v => UArrayMemo v e | v -> e

    -- | Evaluate computation using unboxed array
    --
    -- Key range should cover all possible keys used in computation
    -- otherwise /not in range/ error is generated by array
    evalUArrayMemo :: (Ix k, MArray (UArray m) e m, UArrayMemo v e) =>
                      UArrayCache k e m a  -- ^memoized computation to be evaluated
                   -> (k,k)  -- ^array key range
                   -> m a  -- ^computation result
    {-# INLINE evalUArrayMemo #-}
    evalUArrayMemo :: forall k (m :: * -> *) e v a.
    (Ix k, MArray (UArray m) e m, UArrayMemo v e) =>
    UArrayCache k e m a -> (k, k) -> m a
    evalUArrayMemo = Cache (UArray m) k e m a -> (k, k) -> m a
    forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m a
    genericEvalArrayMemo

    -- | Evaluate computation and the final content of array cache using unboxed array
    --
    -- Key range should cover all possible keys used in computation
    -- otherwise /not in range/ error is generated by array
    runUArrayMemo :: (Ix k, MArray (UArray m) e m, UArrayMemo v e) =>
                     UArrayCache k e m a  -- ^memoized computation to be evaluated
                  -> (k,k)  -- ^array key range
                  -> m (a, UArray m k e)  -- ^computation result and final array cache
    {-# INLINE runUArrayMemo #-}
    runUArrayMemo :: forall k (m :: * -> *) e v a.
    (Ix k, MArray (UArray m) e m, UArrayMemo v e) =>
    UArrayCache k e m a -> (k, k) -> m (a, UArray m k e)
    runUArrayMemo = Cache (UArray m) k e m a -> (k, k) -> m (a, UArray m k e)
    forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m (a, arr k e)
    genericRunArrayMemo

    genericEvalArrayMemo :: (Ix k, MaybeLike e v, MArray arr e m) =>
                            Cache arr k e m a -> (k, k) -> m a
    {-# INLINE genericEvalArrayMemo #-}
    genericEvalArrayMemo :: forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m a
    genericEvalArrayMemo Cache arr k e m a
    m (k, k)
    lu = do
      arr <- (k, k) -> e -> m (arr k e)
    forall i. Ix i => (i, i) -> e -> m (arr i e)
    forall (a :: * -> * -> *) e (m :: * -> *) i.
    (MArray a e m, Ix i) =>
    (i, i) -> e -> m (a i e)
    newArray (k, k)
    lu e
    forall a v. MaybeLike a v => a
    nothing
      evalReaderCache m (Container arr)

    genericRunArrayMemo :: (Ix k, MaybeLike e v, MArray arr e m) =>
                           Cache arr k e m a -> (k, k) -> m (a, arr k e)
    {-# INLINE genericRunArrayMemo #-}
    genericRunArrayMemo :: forall k e v (arr :: * -> * -> *) (m :: * -> *) a.
    (Ix k, MaybeLike e v, MArray arr e m) =>
    Cache arr k e m a -> (k, k) -> m (a, arr k e)
    genericRunArrayMemo Cache arr k e m a
    m (k, k)
    lu = do
      arr <- (k, k) -> e -> m (arr k e)
    forall i. Ix i => (i, i) -> e -> m (arr i e)
    forall (a :: * -> * -> *) e (m :: * -> *) i.
    (MArray a e m, Ix i) =>
    (i, i) -> e -> m (a i e)
    newArray (k, k)
    lu e
    forall a v. MaybeLike a v => a
    nothing
      a <- evalReaderCache m (Container arr)
      return (a, arr)