{-# LANGUAGE GADTs, RankNTypes, KindSignatures, FlexibleInstances, TypeOperators, PolyKinds #-}
module Data.Thrist (
      -- * Types:
        Thrist (..)
      , Flipped (..)
      -- * Fold and map functions:
      , mapThrist
      , foldrThrist
      , foldlThrist
      , foldr1Thrist
      , foldl1Thrist
      -- ** Monadic functions:
      , mapMThrist
      , foldMThrist
      -- * Other list-like functions:
      , appendThrist
      , nullThrist
      , lengthThrist
      ) where

import Prelude hiding ((.), id)
import Control.Category
import Data.Monoid
import Control.Arrow
import Control.Monad


-- | A type-threaded list of binary polymorphic types.
data Thrist :: (k -> k -> *) -> k -> k -> * where
  Nil :: Thrist arr a a
  Cons :: (a `arr` b) -> Thrist arr b c -> Thrist arr a c

instance Monoid (Thrist arr a a) where
  mempty  = Nil
  mappend = appendThrist

instance (Arrow arr) => Arrow (Thrist arr) where
  arr f     = Cons (arr f) Nil
  first Nil = Nil
  first t   = Cons ((foldrThrist (flip (.)) id t) *** id) Nil

instance Category (Thrist arr) where
  id    = Nil
  b . a = appendThrist a b


-- | A newtype wrapper, defined for convenience, that "swaps" the two type
-- variables of a binary type. Can be used to reverse a Thrist using
-- `foldlThrist`. See examples.
newtype Flipped m a b = Flipped { unflip :: m b a }


-- | Equivalent to `foldr` for thrists. Takes a combining function, a value to
-- replace Nil, and a thrist, returning some new binary type.
foldrThrist :: (forall i j . (i `arr` j) -> (j `brr` c) -> (i `brr` c))
            -> (b `brr` c)
            -> Thrist arr a b
            -> (a `brr` c)
foldrThrist _ v Nil        = v
foldrThrist f v (Cons h t) = h `f` (foldrThrist f v t)

-- | Equivalent to (++) for thrists.
appendThrist :: Thrist arr a b -> Thrist arr b c -> Thrist arr a c
appendThrist = flip (foldrThrist Cons)

-- | Equivalent to `map` for thrists. Takes a function from one binary type to
-- another and applies it to each thrist element. For example this could
-- convert a thrist of (a,b) into a thrist of Either a b:
mapThrist :: (forall i j . (i `brr` j) -> (i `arr` j))
          -> Thrist brr a b
          -> Thrist arr a b
mapThrist f = foldrThrist (Cons . f) Nil


-- | Equivalent to `foldl` for `Thrist`s.
foldlThrist :: (forall j k . (a `brr` j) -> (j `arr` k) -> (a `brr` k))
            -> (a `brr` b)
            -> Thrist arr b c
            -> (a `brr` c)
foldlThrist _ v Nil        = v
foldlThrist f v (Cons h t) = foldlThrist f (v `f` h) t



-- | Equivalent to `foldl1` for `Thrist`s.
foldl1Thrist :: (forall i j k. (i `arr` j) -> (j `arr` k) -> (i `arr` k))
             -> Thrist arr a b
             -> (a `arr` b)
foldl1Thrist f (Cons a th) = foldlThrist f a th
foldl1Thrist _ Nil         = error "empty thrist"



-- | Equivalent to `foldr1` for `Thrist`s.
foldr1Thrist :: (forall i j k. (i `arr` j) -> (j `arr` k) -> (i `arr` k))
             -> Thrist arr a b
             -> (a `arr` b)
foldr1Thrist _ (Cons b Nil) = b
foldr1Thrist f (Cons a th)  = f a $ foldr1Thrist f th
foldr1Thrist _ Nil          = error "empty thrist"



-- | Equivalent to `mapM` on `Thrist`s.
mapMThrist :: Monad m =>
             (forall i j . (i `brr` j) -> m (i `arr` j))
           -> Thrist brr a b
           -> m (Thrist arr a b)
mapMThrist _ Nil        = return Nil
mapMThrist f (Cons h t) = liftM2 Cons (f h) (mapMThrist f t)


-- | Equivalent to `foldM` on `Thrist`s.
foldMThrist :: Monad m =>
               (forall j k . (a `brr` j) -> (j `arr` k) -> m (a `brr` k))
            -> (a `brr` b)
            -> Thrist arr b c
            -> m (a `brr` c)
foldMThrist _ a Nil        = return a
foldMThrist f a (Cons h t) = f a h >>= \fah -> foldMThrist f fah t



-- | Returns `True` when the Thrist is `Nil`.
nullThrist :: Thrist arr a b -> Bool
nullThrist Nil = True
nullThrist _   = False

-- | Returns the length of the Thrist.
lengthThrist :: Thrist arr a b -> Int
lengthThrist Nil = 0
lengthThrist (Cons _ rest) = 1 + lengthThrist rest