{-# LANGUAGE CPP #-}
--------------------------------------------------------------------
-- |
-- Module    : Data.Supply
-- Copyright : (c) Iavor S. Diatchki, 2007
-- License   : BSD3
--
-- Maintainer: Iavor S. Diatchki <iavor.diatchki@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.

module Data.Supply
  (

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

  -- * Obtaining values from supplies
  , supplyValue

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

-- Usinga an IORef is thread-safe because we update it with 'atomicModifyIORef'.
import Data.IORef(IORef,newIORef,atomicModifyIORef)

#if __GLASGOW_HASKELL__ >= 608
import GHC.IOBase(unsafeDupableInterleaveIO)
#else
import System.IO.Unsafe(unsafeInterleaveIO)
unsafeDupableInterleaveIO :: IO a -> IO a
unsafeDupableInterleaveIO = unsafeInterleaveIO
#endif

-- 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)


{-# INLINE genericNewSupply #-}
genericNewSupply :: b -> (IORef b -> IO a) -> IO (Supply a)
genericNewSupply start genSym = gen =<< newIORef start

  where gen r = unsafeDupableInterleaveIO
              $ do v <- unsafeDupableInterleaveIO (genSym r)
                   ls <- gen r
                   rs <- gen r
                   return (Node v ls rs)

-- | 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   = genericNewSupply (iterate f x) listGenSym

{-# SPECIALIZE newEnumSupply :: IO (Supply Int) #-}
-- | 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   = genericNewSupply (toEnum 0) enumGenSym

{-# SPECIALIZE newNumSupply :: IO (Supply Int) #-}
-- | 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    = genericNewSupply 0 numGenSym


-- Different ways to generate new values:
listGenSym     :: IORef [a] -> IO a
listGenSym r    = atomicModifyIORef r (\(a:as) -> (as,a))

enumGenSym     :: Enum a => IORef a -> IO a
enumGenSym r    = atomicModifyIORef r (\a -> let n = succ a in seq n (n,a))

numGenSym      :: Num a => IORef a -> IO a
numGenSym r     = atomicModifyIORef r (\a -> let n = 1 + a in seq n (n,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 s f = Node { supplyValue = f s
                        , supplyLeft  = modifySupply (supplyLeft s) f
                        , supplyRight = modifySupply (supplyRight s) f
                        }

-- (Supply, supplyValue, modifySupply) form 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 -----------------------------------------------------------

-- | 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)