Zipping and aligning Traversals

{-# LANGUAGE ScopedTypeVariables, RankNTypes #-}

module ZipTraversal (
    FunList(..),
    runFunList,
    FunApp(..),
    sell,
    runFunApp,
    zipTraversal,
    alignTraversal,
    zipAlignTraversal
  ) where

import Control.Lens
import Control.Applicative (Applicative(..))
import Data.Foldable (Foldable(..))
import Data.These
import Data.Profunctor
import Data.Profunctor.Strong (Costrong(..))
import Data.Profunctor.Choice (Cochoice(..))

data FunList t b a
  = Done t
  | More a (FunList (b -> t) b a)

fmapFL :: forall a a' b t. (a -> a') -> FunList t b a -> FunList t b a'
fmapFL f = go where
  go :: forall u c. FunList u c a -> FunList u c a'
  go (Done t) = Done t
  go (More a fl') = More (f a) (go fl')
  {-# INLINE go #-}
{-# INLINE fmapFL #-}

cmapFL :: forall a a' b t. a' -> FunList t b a -> FunList t b a'
cmapFL a = go where
  go :: forall u c. FunList u c a -> FunList u c a'
  go (Done t) = Done t
  go (More _ fl') = More a (go fl')
  {-# INLINE go #-}
{-# INLINE cmapFL #-}

foldrFL :: forall a b t r. (a -> r -> r) -> r -> FunList t b a -> r
foldrFL f z = go where
  go :: forall u c. FunList u c a -> r
  go (Done _) = z
  go (More a fl') = f a (go fl')
  {-# INLINE go #-}
{-# INLINE foldrFL #-}

foldlFL' :: forall a b t r. (r -> a -> r) -> r -> FunList t b a -> r
foldlFL' f = go where
  go :: forall u c. r -> FunList u c a -> r
  go z _ | z `seq` False = undefined
  go z (Done _) = z
  go z (More a fl') = go (f z a) fl'
  {-# INLINE go #-}
{-# INLINE foldlFL' #-}

lengthWithFL :: Int -> FunList t b a -> Int
lengthWithFL n _ | n `seq` False = undefined
lengthWithFL n (Done _) = n
lengthWithFL n (More _ fl') = lengthWithFL (n + 1) fl'
{-# INLINE lengthWithFL #-}

traverseFL :: forall f a a' b t. Applicative f => (a -> f a') -> FunList t b a -> f (FunList t b a')
traverseFL f = go where
  go :: forall u c. FunList u c a -> f (FunList u c a')
  go (Done t) = pure (Done t)
  go (More a fl') = liftA2 More (f a) (go fl')
  {-# INLINE go #-}
{-# INLINE traverseFL #-}

btmapFL :: forall a b b' t t'. (b' -> b) -> (t -> t') -> FunList t b a -> FunList t' b' a
btmapFL fb = go where
  go :: forall u u' c. (u -> u') -> FunList u b c -> FunList u' b' c
  go ft (Done t) = Done (ft t)
  go ft (More a fl') = More a (go (\bt b -> ft (bt (fb b))) fl')
  {-# INLINE go #-}
{-# INLINE btmapFL #-}

abtmapFL :: forall a a' b b' t t'. (b' -> b) -> (a -> a') -> (t -> t') -> FunList t b a -> FunList t' b' a'
abtmapFL fb fa = go where
  go :: forall u u'. (u -> u') -> FunList u b a -> FunList u' b' a'
  go ft (Done t) = Done (ft t)
  go ft (More a fl') = More (fa a) (go (\bt b -> ft (bt (fb b))) fl')
  {-# INLINE go #-}
{-# INLINE abtmapFL #-}

unfirstFL :: (t -> t') -> FunList t (b, c) (a, c) -> FunList t' b a
unfirstFL ft (Done t) = Done (ft t)
unfirstFL ft (More (a, c) fl') = More a (unfirstFL (\bt b -> ft (bt (b, c))) fl')
{-# INLINE unfirstFL #-}

unsecondFL :: (t -> t') -> FunList t (c, b) (c, a) -> FunList t' b a
unsecondFL ft (Done t) = Done (ft t)
unsecondFL ft (More (c, a) fl') = More a (unsecondFL (\bt b -> ft (bt (c, b))) fl')
{-# INLINE unsecondFL #-}

unleftFL :: (t -> t') -> FunList t (Either b c) (Either a c) -> FunList t' b a
unleftFL ft (Done t) = Done (ft t)
unleftFL ft (More ac fl') = case ac of
  Left a  -> More a (unleftFL (\bt b -> ft (bt (Left b))) fl')
  Right c -> unleftFL (\bt -> ft (bt (Right c))) fl'
{-# INLINE unleftFL #-}

unrightFL :: (t -> t') -> FunList t (Either c b) (Either c a) -> FunList t' b a
unrightFL ft (Done t) = Done (ft t)
unrightFL ft (More ca fl') = case ca of
  Left c  -> unrightFL (\bt -> ft (bt (Left c))) fl'
  Right a -> More a (unrightFL (\bt b -> ft (bt (Right b))) fl')
{-# INLINE unrightFL #-}

instance Functor (FunList t b) where
  fmap = fmapFL
  {-# INLINE fmap #-}

  (<$) = cmapFL
  {-# INLINE (<$) #-}

instance Foldable (FunList t b) where
  foldr = foldrFL
  {-# INLINE foldr #-}

  foldl' = foldlFL'
  {-# INLINE foldl' #-}

  null (Done t) = True
  null (More {}) = False
  {-# INLINE null #-}

  length = lengthWithFL 0
  {-# INLINE length #-}

instance Traversable (FunList t b) where
  traverse = traverseFL
  {-# INLINE traverse #-}

instance Profunctor (FunList t) where
  lmap f = \fl -> case fl of
    Done t -> Done t
    More a fl' -> More a (btmapFL f (. f) fl')
  {-# INLINE lmap #-}

  rmap = fmapFL
  {-# INLINE rmap #-}

  dimap f g = \fl -> case fl of
    Done t -> Done t
    More a fl' -> More (g a) (abtmapFL f g (. f) fl')
  {-# INLINE dimap #-}

instance Costrong (FunList t) where
  unfirst (Done t) = Done t
  unfirst (More (a, c) fl') = More a (unfirstFL (\bt b -> bt (b, c)) fl')
  {-# INLINE unfirst #-}

  unsecond (Done t) = Done t
  unsecond (More (c, a) fl') = More a (unsecondFL (\bt b -> bt (c, b)) fl')
  {-# INLINE unsecond #-}

instance Cochoice (FunList t) where
  unleft (Done t) = Done t
  unleft (More ac fl') = case ac of
    Left a  -> More a (unleftFL (\bt b -> bt (Left b)) fl')
    Right c -> unleftFL (\bt -> bt (Right c)) fl'
  {-# INLINE unleft #-}

  unright (Done t) = Done t
  unright (More ca fl') = case ca of
    Left c  -> unrightFL (\bt -> bt (Left c)) fl'
    Right a -> More a (unrightFL (\bt b -> bt (Right b)) fl')
  {-# INLINE unright #-}

runFunList :: forall f a b t. Applicative f => (a -> f b) -> FunList t b a -> f t
runFunList f = go where
  go :: forall u. FunList u b a -> f u
  go (Done t) = pure t
  go (More a fl') = liftA2 (flip id) (f a) (go fl')
  {-# INLINE go #-}
{-# INLINE runFunList #-}

newtype FunApp a b t = FunApp { unFunApp :: forall r. (forall x. ((t -> r) -> x) -> FunList x b a) -> FunList r b a }

instance Functor (FunApp a b) where
  fmap f (FunApp m) = FunApp $ \y -> m $ \c -> y $ \x -> c $ x . f
  {-# INLINE fmap #-}

  a <$ FunApp m = FunApp $ \y -> m $ \c -> y $ \x -> c $ const $ x a
  {-# INLINE (<$) #-}

instance Applicative (FunApp a b) where
  pure t = FunApp $ \y -> y ($ t)
  {-# INLINE pure #-}

  liftA2 f (FunApp ma) (FunApp mb) = FunApp $ \y -> ma $ \t -> mb $ \u -> y $ \c -> u $ \b -> t $ \a -> c (f a b)
  {-# INLINE liftA2 #-}

  FunApp mf <*> FunApp ma = FunApp $ \y -> mf $ \t -> ma $ \u -> y $ \c -> u $ \a -> t $ \f -> c (f a)
  {-# INLINE (<*>) #-}

  FunApp ma  *> FunApp mb = FunApp $ \y -> ma $ \t -> mb $ \u -> y $ \c -> u $ \b -> t $ const (c b)
  {-# INLINE ( *>) #-}

  FunApp ma <*  FunApp mb = FunApp $ \y -> ma $ \t -> mb $ \u -> y $ \c -> u $ const (t c)
  {-# INLINE (<* ) #-}

sell :: a -> FunApp a b b
sell a = FunApp $ \y -> More a (y id)
{-# INLINE sell #-}

runFunApp :: FunApp a b t -> FunList t b a
runFunApp (FunApp m) = m (\x -> Done (x id))
{-# INLINE runFunApp #-}

zipFunList :: forall f a b c t. Applicative f => ((a, b) -> f c) -> FunList t c a -> FunList t c b -> f t
zipFunList f = go where
  go :: forall u. FunList u c a -> FunList u c b -> f u
  go (Done ta) _ = pure ta
  go _ (Done tb) = pure tb
  go (More a fa') (More b fb') = liftA2 (flip id) (f (a, b)) (go fa' fb')
  {-# INLINE go #-}
{-# INLINE zipFunList #-}

zipTraversal :: LensLike (FunApp a c) sa t a c -> LensLike (FunApp b c) sb t b c -> Traversal (sa, sb) t (a, b) c
zipTraversal ta tb = go where
  go f = \ (sa, sb) -> zipFunList f (runFunApp $ ta sell sa) (runFunApp $ tb sell sb)
  {-# INLINE go #-}
{-# INLINE zipTraversal #-}

alignFunList :: forall f a b c t. Applicative f => (These a b -> f c) -> FunList t c a -> FunList t c b -> f t
alignFunList f = go where
  go :: forall u. FunList u c a -> FunList u c b -> f u
  go fa (Done _) = runFunList (f . This) fa
  go (Done _) fb = runFunList (f . That) fb
  go (More a fa') (More b fb') = liftA2 (flip id) (f (These a b)) (go fa' fb')
  {-# INLINE go #-}
{-# INLINE alignFunList #-}

alignTraversal :: LensLike (FunApp a c) sa t a c -> LensLike (FunApp b c) sb t b c -> Traversal (sa, sb) t (These a b) c
alignTraversal ta tb = go where
  go f = \ (sa, sb) -> alignFunList f (runFunApp $ ta sell sa) (runFunApp $ tb sell sb)
  {-# INLINE go #-}
{-# INLINE alignTraversal #-}

zipAlignFunList :: forall f a b c ta tb. Applicative f => (These a b -> f c) -> FunList ta c a -> FunList tb c b -> f (Ordering, ta, tb)
zipAlignFunList f = go where
  go :: forall ua ub. FunList ua c a -> FunList ub c b -> f (Ordering, ua, ub)
  go (Done ta) (Done tb) = pure (EQ, ta, tb)
  go (Done ta) fb = fmap ((,,) LT ta) (runFunList (f . That) fb)
  go fa (Done tb) = fmap (\ta -> (GT, ta, tb)) (runFunList (f . This) fa)
  go (More a fa') (More b fb') = liftA2 (\c ~(o, ta, tb) -> (o, ta c, tb c)) (f (These a b)) (go fa' fb')
  {-# INLINE go #-}
{-# INLINE zipAlignFunList #-}

zipAlignTraversal :: LensLike (FunApp a c) sa ta a c -> LensLike (FunApp b c) sb tb b c -> Traversal (sa, sb) (Ordering, ta, tb) (These a b) c
zipAlignTraversal ta tb = go where
  go f = \ (sa, sb) -> zipAlignFunList f (runFunApp $ ta sell sa) (runFunApp $ tb sell sb)
  {-# INLINE go #-}
{-# INLINE zipAlignTraversal #-}