{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ExistentialQuantification #-}
module Data.Fold.R
  ( R(..)
  ) where

import Control.Applicative
import Control.Comonad
import Control.Lens
import Data.Foldable hiding (sum, product)
import Data.Fold.Class
import Data.Functor.Extend
import Data.Functor.Bind
import Data.Profunctor.Unsafe
-- import Data.Traversable
import Unsafe.Coerce
import Prelude hiding (foldr, sum, product, length)

-- right folds
data R a b = forall r. R (r -> b) (a -> r -> r) r

-- | leaky 'prefix', efficient 'postfix'
instance Folding R where
  run t (R k h z)     = k (foldr h z t)
  run1 t (R k h z)    = k (h t z)
  runOf l s (R k h z) = k (foldrOf l h z s)
  prefix s            = extend (run s)
  prefix1 a           = extend (run1 a)
  prefixOf l s        = extend (runOf l s)
  postfix t s         = run s (duplicate t)
  postfix1 t a        = run1 a (duplicate t)
  postfixOf l t s     = runOf l s (duplicate t)

instance Profunctor R where
  dimap f g (R k h z) = R (g.k) (h.f) z
  {-# INLINE dimap #-}
  rmap g (R k h z) = R (g.k) h z
  {-# INLINE rmap #-}
  lmap f (R k h z) = R k (h.f) z
  {-# INLINE lmap #-}
  (#.) _ = unsafeCoerce
  {-# INLINE (#.) #-}
  x .# _ = unsafeCoerce x
  {-# INLINE (.#) #-}

instance Choice R where
  left' (R k h z) = R (_Left %~ k) step (Left z) where
    step (Left x) (Left y) = Left (h x y)
    step (Right c) _ = Right c
    step _ (Right c) = Right c
  {-# INLINE left' #-}

  right' (R k h z) = R (_Right %~ k) step (Right z) where
    step (Right x) (Right y) = Right (h x y)
    step (Left c) _ = Left c
    step _ (Left c) = Left c
  {-# INLINE right' #-}

instance Functor (R a) where
  fmap f (R k h z) = R (f.k) h z
  {-# INLINE fmap #-}

  (<$) b = \_ -> pure b
  {-# INLINE (<$) #-}

instance Comonad (R a) where
  extract (R k _ z) = k z
  {-# INLINE extract #-}

  duplicate (R k h z) = R (R k h) h z
  {-# INLINE duplicate #-}

  extend f (R k h z)  = R (f . R k h) h z
  {-# INLINE extend #-}

data Pair a b = Pair !a !b

instance Bind (R a) where
  (>>-) = (>>=)
  {-# INLINE (>>-) #-}

instance Monad (R a) where
  return b = R (\() -> b) (\_ () -> ()) ()
  {-# INLINE return #-}

  m >>= f = R (\xs a -> run xs (f a)) (:) [] <*> m
  {-# INLINE (>>=) #-}

instance Applicative (R a) where
  pure b = R (\() -> b) (\_ () -> ()) ()
  {-# INLINE pure #-}

  R xf bxx xz <*> R ya byy yz = R
    (\(Pair x y) -> xf x $ ya y)
    (\b (Pair x y) -> Pair (bxx b x) (byy b y))
    (Pair xz yz)
  {-# INLINE (<*>) #-}

  (<*) m = \_ -> m
  {-# INLINE (<*) #-}

  _ *> m = m
  {-# INLINE (*>) #-}

instance Extend (R a) where
  extended = extend
  {-# INLINE extended #-}

  duplicated = duplicate
  {-# INLINE duplicated #-}

instance Apply (R a) where
  (<.>) = (<*>)
  {-# INLINE (<.>) #-}

  (<.) m = \_ -> m
  {-# INLINE (<.) #-}

  _ .> m = m
  {-# INLINE (.>) #-}

instance ComonadApply (R a) where
  (<@>) = (<*>)
  {-# INLINE (<@>) #-}

  (<@) m = \_ -> m
  {-# INLINE (<@) #-}

  _ @> m = m
  {-# INLINE (@>) #-}