Some kind of exploration of comonads and zippers in Binary, non-empty trees

ComonadZip.hs

{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE TypeFamilies      #-}

module Bin where

import           Control.Applicative
import qualified Data.Foldable       as F
import           Data.Monoid
import qualified Data.Traversable    as T

import           Test.QuickCheck

{-

Note that there's the concept of the "one-hole context" or data type
derivative which is distinct from the zipper. The type (Non-Empty)
Tree has been designed so that the two align

    (D Tree a, a) == ([D (PF (Tree a)) (Tree a)], Tree a)

This occurs because there's a one-to-one correspondance between
"layers" of the Tree pattern functor (PF) and values embedded in the
tree. As a counterexample, consider the slight variation T = a * T^2 + 1

    data Tree' a = Tree a (Tree a) (Tree a) | Empty

Here, the zipper must allow us to focus on empty leaves while the
one-hole context cannot focus there. Thus, the zipper is

    Z = T + 2aTZ

while the one-hole context and its filler is

    Z = Ca
    Z = aT^2 + 2aTZ

In other words, they're identical excepting that the left branch is
not a full Tree in the case of the one-hole context.

-}

data Tree a =
  Tree  a (Tree a) (Tree a) | Leaf a
  deriving (Eq, Show, Functor, F.Foldable, T.Traversable)

instance Arbitrary a => Arbitrary (Tree a) where
  arbitrary =
    oneof [ Tree <$> arbitrary <*> arbitrary <*> arbitrary
          , Leaf <$> arbitrary
          ]
  shrink (Leaf a) = Leaf <$> shrink a
  shrink (Tree a l r) =
    Tree <$> shrink a <*> shrink l <*> shrink r

data Dir    = L | R                deriving (Eq, Show)

instance Arbitrary Dir where
  arbitrary = fmap (\b -> if b then L else R) arbitrary
  shrink R = [L]
  shrink L = []

data Step a = Step Dir a (Tree a)  deriving (Eq, Show, Functor, F.Foldable, T.Traversable)

instance Arbitrary a => Arbitrary (Step a) where
  arbitrary = Step <$> arbitrary <*> arbitrary <*> arbitrary
  shrink = T.traverse shrink

-- The zipper "list" formulation is more efficiently able to access
-- the current position.
data Zip  a = Zip [Step a] (Tree a) deriving (Eq, Show, Functor)

instance Arbitrary a => Arbitrary (Zip a) where
  arbitrary = Zip <$> arbitrary <*> arbitrary
  shrink (Zip c t) = Zip <$> shrink c <*> shrink t

-- | This folds in the same order as Tree.
-- fold z = fold (load z)
instance F.Foldable Zip where
  foldMap f (Zip cx t) = foldUp cx (F.foldMap f t) where
    foldUp (Step d a sib : cx) rest = case d of
      L -> foldUp cx (f a <> rest <> F.foldMap f sib)
      R -> foldUp cx (f a <> F.foldMap f sib <> rest)

--------------------------------------------------------------------------------

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

instance Comonad Tree where
  extract t = case t of
    Leaf a     -> a
    Tree a _ _ -> a

  duplicate t = case t of
    Leaf _     -> Leaf t
    Tree a l r -> Tree t (duplicate l) (duplicate r)

--------------------------------------------------------------------------------

class (Functor f, Functor (Z f)) => Holey f where
  type Z f :: * -> *
  here :: f a -> Z f a
  load :: Z f a -> f a
  up   :: Z f a -> Either (f a) (Z f a)
  ctxs :: f a -> f (Z f a)

--------------------------------------------------------------------------------

down :: Dir -> Zip a -> Zip a
down d (Zip cs (Tree a l r)) =
  case d of
    L -> Zip (Step L a r : cs) l
    R -> Zip (Step R a l : cs) r

ctxs' :: (Holey f, Comonad (Z f)) => f a -> f (Z f a)
ctxs' = load . duplicate . here

instance Holey Tree where
  type Z Tree = Zip

  -- Tree a -> Zip a
  here = extract . ctxs

  -- Zip a -> Either (Tree a) (Zip a)
  up (Zip [] t) = Left t
  up (Zip (Step d a sib : cs) t) =
    Right . Zip cs $ case d of
      L -> Tree a t sib
      R -> Tree a sib t

  -- Tree a -> Tree (Zip a)
  ctxs t = case t of
    Leaf a     -> Leaf (here t)
    Tree a l r ->
      let z = here t
      in Tree z (down z L l)
                (down z R r)
    where
      down z@(Zip cs (Tree a l r)) d t =
        case d of
          L -> Tree (Zip (Step L a r : cs) l) (down z L l) (down z R r)
          R -> Tree (Zip (Step R a l : cs) r) (down z L l) (down z R r)

  -- Zip a -> Tree a
  load = either id load . up

-- This is pretty ugly, too. Can it be automated?
instance Comonad Zip where
  extract (Zip _ (Tree a _ _)) = a
  duplicate z@(Zip cx (Tree a l r)) =
    Zip (dupCtx z cx) (Tree (id z) (dupT L z l) (dupT R z r))
    where
      dupCtx :: Zip a -> [Step a] -> [Step (Zip a)]
      dupCtx z [] = []
      dupCtx z (Step d a sib : cx) =
        case up z of
          Left (Tree _ _ _) -> []
          Right z' -> Step d (id z') (dupT d z' sib) : dupCtx z' cx

      dupT :: Dir -> Zip a -> Tree t -> Tree (Zip a)
      dupT d z (Tree a l r) =
        let z' = down d z in Tree (id z') (dupT L z' l) (dupT R z' r)

instance Applicative Tree where
  pure a = Tree a (pure a) (pure a)
  Tree f lf rf <*> Tree a la ra = Tree (f a) (lf <*> la) (rf <*> ra)

-- | This would be a bit like Applicative's (<*>) except we can't
-- write a `pure` in infinite, "zippy" form since we must choose a
-- path backwards in history. If we had `Apply` then we could
-- instantiate this there, however.
--
-- The other consequence of this is that we can only "look in the
-- past" if we correctly guess the path!
infixl 4 <@>
(<@>) :: Zip (a -> r) -> (Zip a -> Zip r)
Zip cf f <@> Zip ca a = Zip (zipCx cf ca) (f <*> a) where
  zipCx (Step df f tf : cf) (Step da a ta : ca)
    | df == da  = Step df (f a) (tf <*> ta) : zipCx cf ca
    | otherwise = []

conv -- monoidish, but with a template
     :: a
     -> (a -> a -> (a, a) -> b)
     -- teaches us a new way to fold zippers
     -> (Zip a -> b)
conv d comb z =
  case z of
    Zip [] (Leaf a)                   -> comb d a (d, d)
    Zip (Step _ p _ : _) (Leaf a)     -> comb p a (d, d)
    Zip [] (Tree a l r)               -> comb d a (extract l, extract r)
    Zip (Step _ p _ : _) (Tree a l r) -> comb p a (extract l, extract r)

(&) :: Num a => Zip a -> Zip a -> Sum a
a & b = F.fold ((\a b -> Sum (a * b)) <$> a <@> b)