ifoldlViaIfoldr.hs

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveFoldable      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}

module H6IndexedFolds (module H6IndexedFolds) where

import           Data.Coerce
import           Data.Functor.Compose

data Nat = Z | S Nat

data Vec n a where
    Nil  :: Vec 'Z a
    Cons :: a -> Vec n a -> Vec ('S n) a

deriving instance Foldable (Vec n)

ifoldr :: forall b n a. (forall m. a -> b m -> b ('S m)) -> b 'Z -> Vec n a -> b n
ifoldr f z = go where
    go :: Vec m a -> b m
    go Nil         = z
    go (Cons x xs) = f x $ go xs
{-# INLINE ifoldr #-}

newtype Step b = Step
    { unStep :: forall m. b m -> b ('S m)
    }

newtype IFoldlMotive b n = IFoldlMotive
    { unIFoldlMotive :: forall b'. (forall m. Step b -> b' m -> b' ('S m)) -> b' 'Z -> b' n
    }

ifoldl' :: forall b n a. (forall m. b m -> a -> b ('S m)) -> b 'Z -> Vec n a -> b n
ifoldl' f z0 xs = unIFoldlMotive (ifoldr fm zm xs) unStep z0 where
    fm :: a -> IFoldlMotive b m -> IFoldlMotive b ('S m)
    fm x (IFoldlMotive h) = IFoldlMotive $ \c z ->
        getCompose $ h (coerceC c) $! Compose (c (Step $ \z' -> f z' x) z)
    {-# INLINE fm #-}

    zm = IFoldlMotive $ \_ z -> z
    {-# INLINE zm #-}

    coerceC
        :: (Step b -> b' ('S m) -> b' ('S ('S m)))
        -> Step b -> Compose b' 'S m -> Compose b' 'S ('S m)
    coerceC = coerce
    {-# INLINE coerceC #-}
{-# INLINE ifoldl' #-}

newtype Flip f y x = Flip
    { unFlip :: f x y
    }

reverse :: Vec n a -> Vec n a
reverse = unFlip . ifoldl' (\(Flip xs) x -> Flip $ Cons x xs) (Flip Nil)

-- Here's GHC Core for 'reverse', note how '$WCons' is inlined (there's obviously still overhead
-- associated with the extra argument, but at least we've got inlining sorted out):
{-
Rec {
reverse3
  :: forall {a} {m :: Nat}.
     Vec m a
     -> forall (b' :: Nat -> *).
        (forall (m1 :: Nat). Step (Flip Vec a) -> b' m1 -> b' (S m1))
        -> b' Z -> b' m
reverse3
  = \ (@a_aYw)
      (@(m_aVo :: Nat))
      (ds_d1jx :: Vec m_aVo a_aYw)
      (@(b'_alZ :: Nat -> *))
      (eta_B0
         :: forall (m1 :: Nat).
            Step (Flip Vec a_aYw) -> b'_alZ m1 -> b'_alZ (S m1))
      (eta1_B1 :: b'_alZ Z) ->
      case ds_d1jx of {
        Nil co_aVr -> eta1_B1 `cast` <Co:3> :: ...;
        Cons @n_aVv co_aVw x_amm xs_amn ->
          let! { __DEFAULT ~ nt_s1wa
          <- eta_B0
               ((\ (@(m1_aY3 :: Nat)) (z'_amB :: Flip Vec a_aYw m1_aY3) ->
                   $WCons x_amm (z'_amB `cast` <Co:6> :: ...))
                `cast` <Co:26> :: ...)
               eta1_B1 } in
          (reverse3
             xs_amn
             ((\ (@(m1_aXL :: Nat)) -> eta_B0) `cast` <Co:27> :: ...)
             (nt_s1wa `cast` <Co:7> :: ...))
          `cast` <Co:10> :: ...
      }
end Rec }

reverse2
  :: forall {a} {m :: Nat}.
     Step (Flip Vec a) -> Flip Vec a m -> Flip Vec a (S m)
reverse2
  = \ (@a_aYw)
      (@(m_aYm :: Nat))
      (ds_d1jI :: Step (Flip Vec a_aYw))
      (eta_B0 :: Flip Vec a_aYw m_aYm) ->
      (ds_d1jI `cast` <Co:6> :: ...) eta_B0

reverse1 :: forall {n :: Nat} {a}. Vec n a -> Flip Vec a n
reverse1
  = \ (@(n_aYv :: Nat)) (@a_aYw) (x_X1 :: Vec n_aYv a_aYw) ->
      reverse3 x_X1 reverse2 ($WNil `cast` <Co:7> :: ...)

reverse :: forall (n :: Nat) a. Vec n a -> Vec n a
reverse = reverse1 `cast` <Co:15> :: ...
-}