-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Comonad
-- Copyright   :  2004 Dave Menendez
-- License     :  public domain
-- 
-- Maintainer  :  zednenem@psualum.com
-- Stability   :  experimental
-- Portability :  portable
--
-- This module declares the 'Comonad' class, with instances for
-- 'Id' and @((,) a)@, and defines the 'CoKleisli' arrow.
-----------------------------------------------------------------------------

module Control.Comonad
  (
  -- * The Comonad class
    Comonad(..)
  , (=>>)
  , (.>>)
  , liftW
  
  -- * The coKleisli arrow
  , CoKleisli(..)
  
  -- * The product comonad
  , local
  
  -- * Additional functions
  , sequenceW
  )where

import Control.Arrow
import Control.Functor

infixl 1 =>>, .>>

{-|
There are two ways to define a comonad:

I. Provide definitions for 'fmap', 'extract', and 'duplicate'
satisfying these laws:

> extract . duplicate      == id
> fmap extract . duplicate == id
> duplicate . duplicate    == fmap duplicate . duplicate

II. Provide definitions for 'extract' and 'extend'
satisfying these laws:

> extend extract      == id
> extract . extend f  == f
> extend f . extend g == extend (f . extend g)

('fmap' cannot be defaulted, but a comonad which defines
'extend' may simply set 'fmap' equal to 'liftW'.)

A comonad providing definitions for 'extend' /and/ 'duplicate',
must also satisfy these laws:

> extend f  == fmap f . duplicate
> duplicate == extend id
> fmap f    == extend (f . duplicate)

(The first two are the defaults for 'extend' and 'duplicate',
and the third is the definition of 'liftW'.)
-}

class Functor w => Comonad w where
  extract   :: w a -> a
  duplicate :: w a -> w (w a)
  extend    :: (w a -> b) -> (w a -> w b)
  
  extend f  = fmap f . duplicate
  duplicate = extend id

-- | 'fmap' defined in terms of 'extend'
liftW :: Comonad w => (a -> b) -> (w a -> w b)
liftW f = extend (f . extract)

-- | 'extend' with the arguments swapped. Dual to '>>=' for monads.
(=>>) :: Comonad w => w a -> (w a -> b) -> w b
(=>>) = flip extend

-- | Injects a value into the comonad.
(.>>) :: Comonad w => w a -> b -> w b
w .>> b = extend (\_ -> b) w


--

instance Comonad Id where
  extract (Id x) = x
  duplicate y   = Id y
  extend c w    = Id (c w)

instance Comonad ((,) a) where
  extract   (c,x) = x
  duplicate (c,x) = (c,(c,x))

-- | Calls a comonadic function in a modified context
local :: (c -> c') -> ((c',a) -> a) -> ((c,a) -> a)
local g f (c,x) = f (g c, x)

--

newtype CoKleisli w a b = CoKleisli { unCoKleisli :: w a -> b }

instance Functor (CoKleisli w a) where
  fmap f (CoKleisli g) = CoKleisli (f . g)

instance (Comonad w) => Arrow (CoKleisli w) where
  arr f = CoKleisli (f . extract)

  CoKleisli a >>> CoKleisli b
        = CoKleisli (b . fmap a . duplicate)
  
  CoKleisli a &&& CoKleisli b
        = CoKleisli (a &&& b)
  
  CoKleisli a *** CoKleisli b
        = CoKleisli (a . fmap fst &&& b . fmap snd)
  
  first a  = a *** arr id
  second a = arr id *** a

--

mapW :: Comonad w => (w a -> b) -> w [a] -> [b]
mapW f w | null (extract w) = []
         | otherwise        = f (fmap head w) : mapW f (fmap tail w)

parallelW :: Comonad w => w [a] -> [w a]
parallelW w | null (extract w) = []
            | otherwise        = fmap head w : parallelW (fmap tail w)

unfoldW :: Comonad w => (w b -> (a,b)) -> w b -> [a]
unfoldW f w = fst (f w) : unfoldW f (w =>> snd . f)

-- | Converts a list of comonadic functions into a single function
-- returning a list of values
sequenceW :: Comonad w => [w a -> b] -> w a -> [b]
sequenceW []     w = []
sequenceW (f:fs) w = f w : sequenceW fs w