Comonad.idr

import Data.Vect

infixl 1 =>=

interface Functor w => Comonad (w : Type -> Type) where
  extract : w a -> a

  (=>=) : (w a -> b) -> (w b -> c) -> (w a -> c)
  f =>= g = g . extend f

  extend : (w a -> b) -> w a -> w b
  extend f = map f . duplicate

  duplicate : w a -> w (w a)
  duplicate = extend id

implementation Comonad Stream where
  extract = head
  duplicate stream = stream :: duplicate (tail stream)

data Tree : Type -> Type where
  Leaf : a -> Tree a
  MkTree : (left : Tree a) -> a -> (right : Tree a) -> Tree a

implementation Functor Tree where
  map f (Leaf x) = Leaf (f x)
  map f (MkTree left x right) = MkTree (map f left) (f x) (map f right)

implementation Comonad Tree where
  extract (Leaf x) = x
  extract (MkTree _ x _) = x
  duplicate t@(Leaf x) = Leaf t
  duplicate t@(MkTree left _ right) = MkTree (duplicate left) t (duplicate right)

data FocusVect : Nat -> Type -> Type where
  MkFV : Vect n a -> Fin n -> FocusVect n a

implementation Functor (FocusVect n) where
  map f (MkFV xs n) = MkFV (map f xs) n

implementation Comonad (FocusVect n) where
  extract {n = Z}   (MkFV _ FZ)     impossible
  extract {n = Z}   (MkFV _ (FS _)) impossible
  extract {n = S k} (MkFV (x :: _ ) FZ    ) = x
  extract {n = S k} (MkFV (_ :: xs) (FS m)) = extract (MkFV xs m)
  duplicate (MkFV xs n) = MkFV (map (MkFV xs) range) n

data Store : Type -> Type -> Type where
  MkStore : s -> (s -> a) -> Store s a

implementation Functor (Store s) where
  map f (MkStore s g) = MkStore s (f . g)

implementation Comonad (Store s) where
  extract (MkStore s f) = f s
  duplicate (MkStore s f) = MkStore s (\s' => MkStore s' f)

-- Stream with first 3 elements focused
data Focus3Stream : Type -> Type where
  MkWFS : (a, a, a) -> Stream a -> Focus3Stream a

implementation Functor Focus3Stream where
  map f (MkWFS (h1, h2, h3) t) = MkWFS (f h1, f h2, f h3) (map f t)

focus3 : Stream a -> Focus3Stream a
focus3 (h1 :: h2 :: h3 :: xs) = MkWFS (h1, h2, h3) xs

unfocus : Focus3Stream a -> Stream a
unfocus (MkWFS (h1, h2, h3) xs) = h1 :: h2 :: h3 :: xs

dupFocus3 : Stream a -> Stream (Focus3Stream a)
dupFocus3 xs = focus3 xs :: dupFocus3 (tail xs)

implementation Comonad Focus3Stream where
  extract (MkWFS (h, _, _) _) = h
  duplicate s@(MkWFS (h1, h2, h3) t) = MkWFS (s, focus3 (h2 :: h3 :: t), focus3 (h3 :: t)) (dupFocus3 t)

movingAverage : Focus3Stream Double -> Focus3Stream Double
movingAverage = extend $ \(MkWFS (x, y, z) _) => (x + y + z) / 3

-- *Comonad> take 6 $ unfocus $ movingAverage $ focus3 $ cycle [3.0, 6.0, 0.0, 9.0]
-- [3.0, 5.0, 4.0, 6.0, 3.0, 5.0] : List Double