{-|
Module : Nested
Description : Defines how to nest functors
Maintainer : gatlin@niltag.net

Re-implementation of a library by Kenneth Foner. Upgrades were made to make this
compile and functional dependencies have been replaced by type families.

This is used in 'Sheet' especially in order to automatically build and inspect
multi-dimensional containers out of nested functors.
-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ConstraintKinds #-}


module Nested
  ( -- * Nested functors
    Nested (..)
  , F
  , N
    -- * Deconstructing nested functors
  , UnNest
  , unNest
    -- * Manipulating nested functors
  , NestedCountable(type NestedCount)
  , nestedCount
  , NestedNTimes
  )
where

import Control.Comonad
  ( Comonad(..)
  , ComonadApply(..))
import Control.Applicative (Alternative(..))
import Data.Distributive (Distributive(..))
import Data.Functor.Rep (Representable(..))
import Data.Kind (Type)

import Lists (Counted(..))
import Peano (Natural(..), S, Z)

data F (x :: Type -> Type)
data N (o :: Type) (i :: Type -> Type)

-- | A @Nested fs a@ is the composition of all the layers mentioned in @fs@,
-- applied to an @a@. Specifically, the @fs@ parameter is a sort of snoc-list
-- holding type constructors of kind @(* -> *)@. The outermost layer appears as
-- the parameter to @Flat@; the innermost layer appears as the rightmost
-- argument to the outermost @Nest@.
data Nested fs a
  = forall f. (fs ~ F f) => Flat (f a)
  | forall fs' f. (fs ~ N fs' f) => Nest (Nested fs' (f a))

type family UnNest x where
  UnNest (Nested (F f) a) = f a
  UnNest (Nested (N fs f) a) = Nested fs (f a)

unNest :: Nested fs a -> UnNest (Nested fs a)
unNest :: forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest (Flat !f a
x) = f a
x
unNest (Nest !Nested fs' (f a)
x) = Nested fs' (f a)
x

instance Functor f => Functor (Nested (F f)) where
  fmap :: forall a b. (a -> b) -> Nested (F f) a -> Nested (F f) b
fmap a -> b
f !Nested (F f) a
x = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat f b
mapped where
    !x' :: UnNest (Nested (F f) a)
x' = forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest Nested (F f) a
x
    !mapped :: f b
mapped = let fx :: f b
fx = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f UnNest (Nested (F f) a)
x' in f b
fx

instance (Functor f, Functor (Nested fs)) => Functor (Nested (N fs f)) where
  fmap :: forall a b. (a -> b) -> Nested (N fs f) a -> Nested (N fs f) b
fmap a -> b
f !Nested (N fs f) a
x = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest Nested fs (f b)
mapped where
    !x' :: UnNest (Nested (N fs f) a)
x' = forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest Nested (N fs f) a
x
    !mapped :: Nested fs (f b)
mapped = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap f a -> f b
f' UnNest (Nested (N fs f) a)
x'
    f' :: f a -> f b
f' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f

instance (Applicative f) => Applicative (Nested (F f)) where
  pure :: forall a. a -> Nested (F f) a
pure = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Flat f (a -> b)
f <*> :: forall a b.
Nested (F f) (a -> b) -> Nested (F f) a -> Nested (F f) b
<*> Flat f a
x = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat (f (a -> b)
f forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f a
x)

instance ( Applicative f
         , Applicative (Nested fs))
  => Applicative (Nested (N fs f)) where
  pure :: forall a. a -> Nested (N fs f) a
pure = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Nest Nested fs' (f (a -> b))
f <*> :: forall a b.
Nested (N fs f) (a -> b) -> Nested (N fs f) a -> Nested (N fs f) b
<*> Nest Nested fs' (f a)
x = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest (forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Nested fs' (f (a -> b))
f forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Nested fs' (f a)
x)

instance (Comonad f) => Comonad (Nested (F f)) where
  extract :: forall a. Nested (F f) a -> a
extract = forall (w :: * -> *) a. Comonad w => w a -> a
extract forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest
  duplicate :: forall a. Nested (F f) a -> Nested (F f) (Nested (F f) a)
duplicate = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance ( Comonad f
         , Comonad (Nested fs)
         , Functor (Nested (N fs f))
         , Distributive f )
  => Comonad (Nested (N fs f)) where
  extract :: forall a. Nested (N fs f) a -> a
extract = forall (w :: * -> *) a. Comonad w => w a -> a
extract forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (w :: * -> *) a. Comonad w => w a -> a
extract forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest
  duplicate :: forall a. Nested (N fs f) a -> Nested (N fs f) (Nested (N fs f) a)
duplicate =
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest        -- wrap it again
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute       -- swap middle two layers
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate             -- duplicate outer functor
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (w :: * -> *) a. Comonad w => w a -> w (w a)
duplicate        -- duplicate inner functor
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest                -- note: can't pattern match

instance (ComonadApply f) => ComonadApply (Nested (F f)) where
  Flat f (a -> b)
f <@> :: forall a b.
Nested (F f) (a -> b) -> Nested (F f) a -> Nested (F f) b
<@> Flat f a
x = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat (f (a -> b)
f forall (w :: * -> *) a b.
ComonadApply w =>
w (a -> b) -> w a -> w b
<@> f a
x)

instance ( ComonadApply f
         , Distributive f
         , ComonadApply (Nested fs))
  => ComonadApply (Nested (N fs f)) where
  Nest Nested fs' (f (a -> b))
f <@> :: forall a b.
Nested (N fs f) (a -> b) -> Nested (N fs f) a -> Nested (N fs f) b
<@> Nest Nested fs' (f a)
x = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest (forall (w :: * -> *) a b.
ComonadApply w =>
w (a -> b) -> w a -> w b
(<@>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Nested fs' (f (a -> b))
f forall (w :: * -> *) a b.
ComonadApply w =>
w (a -> b) -> w a -> w b
<@> Nested fs' (f a)
x)

instance (Distributive f) => Distributive (Nested (F f)) where
  distribute :: forall (f :: * -> *) a.
Functor f =>
f (Nested (F f) a) -> Nested (F f) (f a)
distribute = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance ( Distributive f
         , Distributive (Nested fs))
  => Distributive (Nested (N fs f)) where
  distribute :: forall (f :: * -> *) a.
Functor f =>
f (Nested (N fs f) a) -> Nested (N fs f) (f a)
distribute = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a.
(Distributive g, Functor f) =>
f (g a) -> g (f a)
distribute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance (Representable f) => Representable (Nested (F f)) where
  type Rep (Nested (F f)) = Counted (S Z) (Rep f)
  index :: forall a. Nested (F f) a -> Rep (Nested (F f)) -> a
index !Nested (F f) a
obj (Rep f
key ::: Counted t (Rep f)
CountedNil) =
    let !obj' :: UnNest (Nested (F f) a)
obj' = forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest Nested (F f) a
obj 
        !r :: a
r = forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index UnNest (Nested (F f) a)
obj' Rep f
key
    in  a
r
  tabulate :: forall a. (Rep (Nested (F f)) -> a) -> Nested (F f) a
tabulate Rep (Nested (F f)) -> a
describe = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> a
describe' where
    describe' :: Rep f -> a
describe' !Rep f
key = let !key' :: Counted (S Z) (Rep f)
key' = Rep f
key forall n a t. (n ~ S t) => a -> Counted t a -> Counted n a
::: forall n a. (n ~ Z) => Counted n a
CountedNil in Rep (Nested (F f)) -> a
describe Counted (S Z) (Rep f)
key'

instance ( Representable f
         , fs ~ NestedNTimes (NestedCount fs) f
         , Representable (Nested fs)
         , Rep (Nested fs) ~ Counted (NestedCount fs) (Rep f))
  => Representable (Nested (N fs f)) where
  type Rep (Nested (N fs f)) = Counted (S (NestedCount fs)) (Rep f)
  index :: forall a. Nested (N fs f) a -> Rep (Nested (N fs f)) -> a
index !Nested (N fs f) a
obj (!Rep f
k ::: !Counted t (Rep f)
ks) =
    let !obj' :: UnNest (Nested (N fs f) a)
obj' = forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest Nested (N fs f) a
obj
        !mapped :: Nested fs a
mapped = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(!f a
o) -> forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index f a
o Rep f
k) UnNest (Nested (N fs f) a)
obj'
        !v :: a
v = forall (f :: * -> *) a. Representable f => f a -> Rep f -> a
index Nested fs a
mapped Counted t (Rep f)
ks
    in  a
v
  tabulate :: forall a. (Rep (Nested (N fs f)) -> a) -> Nested (N fs f) a
tabulate Rep (Nested (N fs f)) -> a
describe = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest Nested fs (f a)
tabulated where
    !tabulated :: Nested fs (f a)
tabulated = forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Counted (NestedCount fs) (Rep f) -> f a
fn
    fn :: Counted (NestedCount fs) (Rep f) -> f a
fn !Counted (NestedCount fs) (Rep f)
ks = forall (f :: * -> *) a. Representable f => (Rep f -> a) -> f a
tabulate Rep f -> a
fn' where
      fn' :: Rep f -> a
fn' !Rep f
k = let !x :: a
x = Rep (Nested (N fs f)) -> a
describe (Rep f
k forall n a t. (n ~ S t) => a -> Counted t a -> Counted n a
::: Counted (NestedCount fs) (Rep f)
ks) in a
x

instance (Foldable f) => Foldable (Nested (F f)) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Nested (F f) a -> m
foldMap a -> m
f = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance ( Foldable f
         , Foldable (Nested fs))
  => Foldable (Nested (N fs f)) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Nested (N fs f) a -> m
foldMap a -> m
f = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance (Traversable f) => Traversable (Nested (F f)) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Nested (F f) a -> f (Nested (F f) b)
traverse a -> f b
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance ( Traversable f
         , Traversable (Nested fs))
  => Traversable (Nested (N fs f)) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Nested (N fs f) a -> f (Nested (N fs f) b)
traverse a -> f b
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fs a. Nested fs a -> UnNest (Nested fs a)
unNest

instance (Alternative f) => Alternative (Nested (F f)) where
  empty :: forall a. Nested (F f) a
empty = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat forall (f :: * -> *) a. Alternative f => f a
empty
  Flat f a
x <|> :: forall a. Nested (F f) a -> Nested (F f) a -> Nested (F f) a
<|> Flat f a
y = forall fs a (f :: * -> *). (fs ~ F f) => f a -> Nested fs a
Flat (f a
x forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> f a
y)

instance ( Applicative f
         , Alternative (Nested fs))
  => Alternative (Nested (N fs f)) where
  empty :: forall a. Nested (N fs f) a
empty = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest forall (f :: * -> *) a. Alternative f => f a
empty
  Nest Nested fs' (f a)
x <|> :: forall a.
Nested (N fs f) a -> Nested (N fs f) a -> Nested (N fs f) a
<|> Nest Nested fs' (f a)
y = forall fs a fs' (f :: * -> *).
(fs ~ N fs' f) =>
Nested fs' (f a) -> Nested fs a
Nest (Nested fs' (f a)
x forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Nested fs' (f a)
y)

class NestedCountable (x :: k) where
  type NestedCount x :: Type

instance NestedCountable (F f) where
  type NestedCount (F f) = S Z

instance NestedCountable (N fs f) where
  type NestedCount (N fs f) = S (NestedCount fs)

nestedCount :: Nested fs a -> Natural (NestedCount fs)
nestedCount :: forall fs a. Nested fs a -> Natural (NestedCount fs)
nestedCount (Flat f a
_) = forall n t. (n ~ S t) => Natural t -> Natural n
Succ forall n. (n ~ Z) => Natural n
Zero
nestedCount (Nest Nested fs' (f a)
x) = forall n t. (n ~ S t) => Natural t -> Natural n
Succ (forall fs a. Nested fs a -> Natural (NestedCount fs)
nestedCount Nested fs' (f a)
x)

type family NestedNTimes (n :: Type) (f :: Type -> Type) where
  NestedNTimes (S Z) f = F f
  NestedNTimes (S n) f = N (NestedNTimes n f) f