--------------------------------------------------------------------
-- |
-- Module    : Control.Comonad.Supply
-- Copyright : (c) Edward Kmett 2008
--             (c) Iavor S. Diatchki, 2007
-- License   : BSD3
--
-- Maintainer: Edward Kmett <ekmett@gmail.com>
-- Stability : provisional
-- Portability: portable
--
-- The technique for generating new values is based on the paper
-- ''On Generating Unique Names''
-- by Lennart Augustsson, Mikael Rittri, and Dan Synek.
-- 
-- Integrated from value-supply-0.1
--
-- TODO: a SupplyT Comonad Transformer
--------------------------------------------------------------------

module Control.Comonad.Supply
  ( module Control.Comonad


  -- * Creating supplies
  , Supply
  , newSupply
  , newEnumSupply
  , newNumSupply

  -- * Obtaining values from supplies
  , supplyValue

  -- * Generating new supplies from old
  , supplyLeft
  , supplyRight
  , modifySupply
  , split
  , split2
  , split3
  , split4
  ) where

import Control.Comonad
-- Using 'MVar's might be a bit heavy but it ensures that
-- multiple threads that share a supply will get distinct names.
import Control.Concurrent.MVar
import Control.Functor.Extras
import System.IO.Unsafe(unsafePerformIO)

-- Basics ----------------------------------------------------------------------

-- | A type that can be used to generate values on demand.
-- A supply may be turned into two different supplies by using
-- the functions 'supplyLeft' and 'supplyRight'.
data Supply a = Node
  { -- | Get the value of a supply.  This function, together with
    -- 'modifySupply' forms a comonad on 'Supply'.
    supplyValue :: a

  -- | Generate a new supply.  This supply is different from the one
  -- generated with 'supplyRight'.
  , supplyLeft  :: Supply a

  -- | Generate a new supply. This supply is different from the one
  -- generated with 'supplyLeft'.
  , supplyRight :: Supply a
  }

instance Functor Supply where
  fmap f s = modifySupply s (f . supplyValue)

-- | Creates a new supply of values.
-- The arguments specify how to generate values:
-- the first argument is an initial value, the
-- second specifies how to generate a new value from an existing one.
newSupply    :: a -> (a -> a) -> IO (Supply a)
newSupply x f = fmap (gen True) (newMVar (iterate f x))

  -- The extra argument to ``gen'' is passed because without
  -- it Hugs spots that the recursive calls are the same but does
  -- not know that unsafePerformIO is unsafe.
  where gen _ r = Node { supplyValue  = unsafePerformIO (genSym r),
                         supplyLeft   = gen False r,
                         supplyRight  = gen True r }

        genSym       :: MVar [a] -> IO a
        genSym r      = do a : as <- takeMVar r
                           putMVar r as
                           return a

-- | Generate a new supply by systematically applying a function
-- to an existing supply.  This function, together with 'supplyValue'
-- form a comonad on 'Supply'.
modifySupply :: Supply a -> (Supply a -> b) -> Supply b
modifySupply = flip extend

-- (Supply, supplyValue, modifySupply) forms a comonad:
{-
law1 s      = [ modifySupply s supplyValue, s ]
law2 s f    = [ supplyValue (modifySupply s f), f s ]
law3 s f g  = [ (s `modifySupply` f) `modifySupply` g
              ,  s `modifySupply` \s1 -> g (s1 `modifySupply` f)
              ]
-}


-- Derived functions -----------------------------------------------------------

-- | A supply of values that are in the 'Enum' class.
-- The initial value is @toEnum 0@, new values are generates with 'succ'.
newEnumSupply  :: (Enum a) => IO (Supply a)
newEnumSupply   = newSupply (toEnum 0) succ

-- | A supply of values that are in the 'Num' class.
-- The initial value is 0, new values are generated by adding 1.
newNumSupply   :: (Num a) => IO (Supply a)
newNumSupply    = newSupply 0 (1+)

-- | Generate an infinite list of supplies by using 'supplyLeft' and
-- 'supplyRight' repeatedly.
split          :: Supply a -> [Supply a]
split s         = supplyLeft s : split (supplyRight s)

-- | Split a supply into two different supplies.
-- The resulting supplies are different from the input supply.
split2         :: Supply a -> (Supply a, Supply a)
split2 s        = (supplyLeft s, supplyRight s)

-- | Split a supply into three different supplies.
split3         :: Supply a -> (Supply a, Supply a, Supply a)
split3 s        = let s1 : s2 : s3 : _ = split s
                  in (s1,s2,s3)

-- | Split a supply into four different supplies.
split4         :: Supply a -> (Supply a, Supply a, Supply a, Supply a)
split4 s        = let s1 : s2 : s3 : s4 : _ = split s
                  in (s1,s2,s3,s4)

instance Copointed Supply where
    extract = supplyValue

instance Comonad Supply where
    extend f s = Node { supplyValue = f s
                      , supplyLeft  = modifySupply (supplyLeft s) f
                      , supplyRight = modifySupply (supplyRight s) f
                      }

instance FunctorSplit Supply where
    fsplit = split2