Type indexed foldl' from type-indexed foldr

IxFoldl.hs

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE DeriveTraversable      #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

module IxFoldl
    ( Nat (..)
    , Vec (..)
    , ifoldl'
    ) where
import Data.Functor.Compose
import Data.Type.Equality
import Data.Kind (Type)
import Data.Functor.Product
import Data.Tagged
import Data.Proxy

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 #-}

-- | Addition, but *not* the usual definition. This one shifts
-- succs from the first nat to the second till none are left.
type Plus :: Nat -> Nat -> Nat
type family a `Plus` b where
  'Z `Plus` n = n
  'S m `Plus` n = m `Plus` 'S n

data Fish b jump = Fish
  { getFish :: forall m. b m -> b (jump `Plus` m)
  , pf1 :: !(jump `Plus` 'Z  :~: jump)
  , pf2 :: !(forall m. Tagged m ('S (jump `Plus` m) :~: jump `Plus` 'S m)) }

bumpProxy :: proxy n -> Proxy ('S n)
bumpProxy _ = Proxy

-- | Caution: This walks the whole vector before producing anything. The trouble
-- is that we need to construct a proof that n `Plus` 'Z :~: n, which requires
-- walking everything.
ifoldl' :: forall b a n. (forall m. b m -> a -> b ('S m)) -> b 'Z -> Vec n a -> b n
ifoldl' f z = \vec -> case ifoldr go (Fish id Refl (Tagged Refl)) vec of
   Fish fun Refl _ -> fun z
  where
   go :: forall m. a -> Fish b m -> Fish b ('S m)
   go a (Fish q Refl pf) = Fish (\ !s -> q (f s a))
     (case proxy pf (Proxy @'Z) of Refl -> Refl)
     (unproxy $ \px -> proxy pf (bumpProxy px))

deriving instance Functor (Vec n)
deriving instance Traversable (Vec n)
deriving instance Eq a => Eq (Vec n a)
deriving instance Ord a => Ord (Vec n a)
deriving instance Show a => Show (Vec n a)