{-# LANGUAGE Rank2Types #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Free
-- Copyright   :  2004 Dave Menendez
-- License     :  BSD3
-- 
-- Maintainer  :  dan.doel@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- An implementation of the cofree comonad of a functor, used in
-- histomorphisms and chronomorphisms in Control.Recursion. The
-- cofree comonad can also be seen as a stream parameterized by a
-- functor that controls its branching factor.
--
-----------------------------------------------------------------------------

module Control.Comonad.Cofree
  ( Cofree(..)
  , headCofree
  , tailCofree
  , anaCofree
  , cofreeToList
  , distribCofree
  ) where

import Control.Arrow ((&&&),(***),(>>>), second)
import Control.Comonad

{-|
The cofree comonad of a functor @h@ (also known as an H-branching stream).
Various comonads are a special instance of the cofree comonad:

* @Cofree Identity@ is an infinite stream

* @Cofree Maybe@ is a non-empty stream

* @Cofree []@ is a rose tree

formally:

> Cofree H A = nu X. A * HX
-}
data Cofree h a = Cofree { unCofree :: (a, h (Cofree h a)) }

-- | anamorphism for building a cofree comonad from a seed
anaCofree :: Functor h => (a -> b) -> (a -> h a) -> a -> Cofree h b
anaCofree g1 g2 = g1 &&& fmap (anaCofree g1 g2) . g2 >>> Cofree

headCofree :: Cofree h a -> a
headCofree = fst . unCofree

tailCofree :: Cofree h a -> h (Cofree h a)
tailCofree = snd . unCofree

instance Functor h => Functor (Cofree h) where
  fmap g = unCofree >>> g *** fmap (fmap g) >>> Cofree

instance Functor h => Comonad (Cofree h) where
  extract   = headCofree
  duplicate = anaCofree id tailCofree

-- | Converts a value of the cofree comonad over Maybe into a non-empty list.
cofreeToList :: Cofree Maybe a -> [a]
cofreeToList = unCofree >>> second (maybe [] cofreeToList) >>> uncurry (:) 

-- | Lifts a distributive law of @f@ over @h@ to a distributive law
-- of @f@ over @Cofree h@.
distribCofree :: (Functor h, Functor f) =>
                 (forall a. f (h a) -> h (f a))
                   -> (forall a. f (Cofree h a) -> Cofree h (f a))
distribCofree d = anaCofree (fmap headCofree) (d . fmap tailCofree)