{-# LANGUAGE TypeOperators, TypeFamilies, TupleSections #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}

{-# OPTIONS_GHC -Wall #-}

----------------------------------------------------------------------
-- |
-- Module      :  LambdaCCC.StateTrie
-- Copyright   :  (c) 2014 Tabula, Inc.
-- 
-- Maintainer  :  conal@tabula.com
-- Stability   :  experimental
-- 
-- Memoizing State monad
----------------------------------------------------------------------

module LambdaCCC.StateTrie
  ( StateTrieX, StateTrie(..)
  , toState, fromState
  , get, put, runStateTrie, evalStateTrie, execStateTrie
  ) where

import Control.Arrow (first)
import Control.Applicative (Applicative(..))

import Control.Monad.State -- mtl

import Data.MemoTrie (HasTrie(..),(:->:))

import Circat.Rep (Rep,HasRep(..))

-- | 'StateTrie' inner representation
type StateTrieX s a = s :->: (a,s)

-- | Memoizing state monad
newtype StateTrie s a = StateTrie { unStateTrie :: StateTrieX s a }

-- | Operate inside a 'StateTrie'.
inStateTrie :: (StateTrieX s a -> StateTrieX t b)
            -> (StateTrie  s a -> StateTrie  t b)
inStateTrie = StateTrie <~ unStateTrie
{-# INLINE inStateTrie #-}

{- unused

inStateTrie2 :: (StateTrieX s a -> StateTrieX t b -> StateTrieX u c)
             -> (StateTrie  s a -> StateTrie  t b -> StateTrie  u c)
inStateTrie2 = inStateTrie <~ unStateTrie

-}

-- | Run a memoized stateful computation
runStateTrie :: HasTrie s => StateTrie s a -> s -> (a,s)
runStateTrie (StateTrie t) = untrie t
{-# INLINE runStateTrie #-}

-- | Run a memoized stateful computation and return just value
evalStateTrie :: HasTrie s => StateTrie s a -> s -> a
evalStateTrie = (result.result) fst runStateTrie
{-# INLINE evalStateTrie #-}

-- | Run a memoized stateful computation and return just state
execStateTrie :: HasTrie s => StateTrie s a -> s -> s
execStateTrie = (result.result) snd runStateTrie
{-# INLINE execStateTrie #-}

instance HasTrie s => Functor (StateTrie s) where
  fmap = inStateTrie . fmap . first
  {-# INLINE fmap #-}

instance HasTrie s => Applicative (StateTrie s) where
  pure a = StateTrie (trie (a,))
  (<*>)  = inState2 (<*>)
  {-# INLINE pure #-}
  {-# INLINE (<*>) #-}

-- | 'State'-to-'StateTrie' adapter
fromState :: HasTrie s => State s a -> StateTrie s a
fromState = StateTrie . trie . runState
{-# INLINE fromState #-}

-- | 'StateTrie'-to-'State' adapter
toState :: HasTrie s => StateTrie s a -> State s a
toState = state . untrie . unStateTrie
{-# INLINE toState #-}

-- | Transform using 'State' view
inState :: (HasTrie s, HasTrie t) =>
           (State     s a -> State     t b)
        -> (StateTrie s a -> StateTrie t b)
inState = fromState <~ toState
{-# INLINE inState #-}

-- | Transform using 'State' view
inState2 :: (HasTrie s, HasTrie t, HasTrie u) =>
            (State     s a -> State     t b -> State     u c)
         -> (StateTrie s a -> StateTrie t b -> StateTrie u c)
inState2 = inState <~ toState
{-# INLINE inState2 #-}

instance HasTrie s => Monad (StateTrie s) where
  return  = pure
  m >>= f = joinST (fmap f m)
  {-# INLINE return #-}
  {-# INLINE (>>=) #-}

joinST :: HasTrie s => StateTrie s (StateTrie s a) -> StateTrie s a
joinST = fromState . join . fmap toState . toState
{-# INLINE joinST #-}

-- joinST = inState (join . fmap toState)
--        = inState ((=<<) toState)

instance HasTrie s => MonadState s (StateTrie s) where
  state = StateTrie . trie

-- TODO: Perhaps use 'state' in the definitions of pure and fromState.

type instance Rep (StateTrie s a) = StateTrieX s a
instance HasRep (StateTrie s a) where
  repr (StateTrie t) = t
  abst = StateTrie


{--------------------------------------------------------------------
    Misc
--------------------------------------------------------------------}

-- | Add post- & pre-processing
(<~) :: (b -> b') -> (a' -> a) -> ((a -> b) -> (a' -> b'))
(h <~ f) g = h . g . f

-- | Add post-processing
result :: (b -> b') -> ((a -> b) -> (a -> b'))
result = (.)
-- result = (<~ id)
