{-# LANGUAGE Haskell2010 #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Control.Monad.Trans.Fresh.Delta where

import Control.Applicative
import Control.Arrow (first, second)
import Control.Comonad.Cofree
import Control.Monad.Trans.Class
import Data.Bits
import Data.Fresh
import Data.Triplet

newtype FreshT v m a = FreshT
    { runFreshT :: Cofree Triplet v -> m (a, Cofree Triplet v) }

-- | Creates a @Cofree Triplet v@, based on a seed, and functions to split
--   and succeed it.
seedDelta
    :: v                -- ^ Delta seed
    -> (v -> v)         -- ^ Successor
    -> (v -> (v, v))    -- ^ Splitter
    -> Cofree Triplet v
seedDelta seed sukk splt = go seed
  where
    go tt = tt `seq` (tt :< Triplet (go ll) (go mm) (go rr))
      where
        mm = sukk tt
        (ll, rr) = splt tt

-- | A delta of integer pairs.
integerPairDelta :: (Integer, Integer) -> Cofree Triplet (Integer, Integer)
integerPairDelta n = seedDelta n sukk splt
  where
    sukk = second succ
    splt (m, _) = let m' = m `shiftL` 1
                  in ((m', 0), (m' .|. 1, 0))

instance Functor m => Functor (FreshT v m) where
    fmap f (FreshT xx) = FreshT $ \s -> fmap (first f) (xx s)

instance Applicative m => Applicative (FreshT v m) where
    pure x = FreshT $ \s -> pure (x, s)
    FreshT ff <*> FreshT xx = FreshT $ \(_ :< Triplet l _ r)
        -> (\(f, sf) (x, _) -> (f x, sf)) <$> ff l <*> xx r

instance Monad m => Monad (FreshT v m) where
    return x = FreshT $ \s -> return (x, s)
    FreshT aa >>= f = FreshT $ \(_ :< Triplet l _ r) -> do
        (a, _) <- aa l
        let FreshT ff = f a
        ff r

-- TODO Alternative, MonadPlus
-- TODO MonadFix

instance MonadTrans (FreshT v) where
    lift m = FreshT $ \s -> do
                r <- m
                return (r, s)

-- TODO MonadIO

instance Applicative m => Fresh v (FreshT v m) where
    fresh = FreshT $ \(s :< Triplet _ s' _) -> pure (s, s')