{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, Rank2Types,
             FlexibleInstances, FlexibleContexts, UndecidableInstances #-}
--
-- Experimental interface for returning multiple arrays from an ST calculation
--
-- Author:      Bertram Felgenhauer
-- Date:        2008-03-13
-- Tested with: ghc 6.8.2
--
module EvalST (EvalST (..), evalST, STPure, STArray', STUArray', STPair) where

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Array
import Data.Array.Base
import Data.Array.Unboxed

class EvalST st pure | st -> pure where
    freezeST :: st s -> ST s pure

evalST :: EvalST st pure => (forall s. ST s (st s)) -> pure
evalST f = runST (f >>= freezeST)

-- Basic wrappers
newtype STPure a s = STPure { unSTPure :: a }
newtype STArray' i e s = STArray' { unSTArray' :: STArray s i e }
newtype STUArray' i e s = STUArray' { unSTUArray' :: STUArray s i e }

-- Basic instances
instance EvalST (STPure a) a where
    freezeST = return . unSTPure

instance (Ix i) => EvalST (STArray' i e) (Array i e) where
    freezeST = unsafeFreeze . unSTArray'

instance (IArray UArray e, Ix i) => EvalST (STUArray' i e) (UArray i e) where
    freezeST = unsafeFreezeSTUArray . unSTUArray'

-- Complex instances
data STPair l r s = STPair (l s) (r s)

instance (EvalST st1 pure1, EvalST st2 pure2) => EvalST (STPair st1 st2) (pure1, pure2) where
    freezeST = \(STPair l r) -> liftM2 (,) (freezeST l) (freezeST r)