content: {-# LANGUAGE ConstraintKinds       #-}
  {-# LANGUAGE DataKinds             #-}
  {-# LANGUAGE DeriveDataTypeable    #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PolyKinds             #-}
  {-# LANGUAGE RankNTypes            #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE TypeFamilies          #-}
  {-# LANGUAGE TypeOperators         #-}
  {-# LANGUAGE UndecidableInstances  #-}
  module T21689a
    ( Arity
    , Dim
    , DimM
    , IVector
    , Mutable
    , MVector
    , Vector(..)
    , constructVec
    , inspectVec
    , gfoldl'
    , gunfold'
    ) where
  
  import Control.Monad.ST (ST, runST)
  import Data.Coerce (coerce)
  import Data.Data (Data)
  import Data.Functor.Const (Const(..))
  import Data.Kind (Type)
  import GHC.TypeLits (KnownNat, Nat, type (+), type (-))
  
  -----
  -- Data.Vector.Fixed.Cont
  -----
  
  data PeanoNum = Z
                | S PeanoNum
  
  type family Peano (n :: Nat) :: PeanoNum where
    Peano 0 = 'Z
    Peano n = 'S (Peano (n - 1))
  
  type family Fn (n :: PeanoNum) (a :: Type) (b :: Type) where
    Fn 'Z     a b = b
    Fn ('S n) a b = a -> Fn n a b
  
  newtype Fun n a b = Fun { unFun :: Fn n a b }
  
  type family Dim (v :: Type -> Type) :: Nat
  
  class Arity (Dim v) => Vector v a where
    construct :: Fun (Peano (Dim v)) a (v a)
  
    inspect   :: v a -> Fun (Peano (Dim v)) a b -> b
  
  type Arity n = ( ArityPeano (Peano n)
                 , KnownNat n
                 , Peano (n+1) ~ 'S (Peano n)
                 )
  
  class ArityPeano n where
    accum :: (forall k. t ('S k) -> a -> t k)
          -> (t 'Z -> b)
          -> t n
          -> Fun n a b
  
    applyFun :: (forall k. t ('S k) -> (a, t k))
             -> t n
             -> (CVecPeano n a, t 'Z)
  
    gunfoldF :: (Data a)
             => (forall b x. Data b => c (b -> x) -> c x)
             -> T_gunfold c r a n -> c r
  
  newtype T_gunfold c r a n = T_gunfold (c (Fn n a r))
  
  gfoldl' :: forall c v a. (Vector v a, Data a)
          => (forall x y. Data x => c (x -> y) -> x -> c y)
          -> (forall x  . x -> c x)
          -> v a -> c (v a)
  gfoldl' f inj v
    = inspect v
    $ gfoldlF f (inj $ unFun (construct :: Fun (Peano (Dim v)) a (v a)))
  
  gunfold' :: forall con c v a. (Vector v a, Data a)
           => (forall b r. Data b => c (b -> r) -> c r)
           -> (forall r. r -> c r)
           -> con -> c (v a)
  gunfold' f inj _
    = gunfoldF f gun
    where
      con = construct                   :: Fun (Peano (Dim v)) a (v a)
      gun = T_gunfold (inj $ unFun con) :: T_gunfold c (v a) a (Peano (Dim v))
  
  gfoldlF :: (ArityPeano n, Data a)
          => (forall x y. Data x => c (x -> y) -> x -> c y)
          -> c (Fn n a r) -> Fun n a (c r)
  gfoldlF f c0 = accum
    (\(T_gfoldl c) x -> T_gfoldl (f c x))
    (\(T_gfoldl c)   -> c)
    (T_gfoldl   c0)
  
  newtype T_gfoldl c r a n = T_gfoldl (c (Fn n a r))
  
  newtype ContVec n a = ContVec (forall r. Fun (Peano n) a r -> r)
  
  type instance Dim (ContVec n) = n
  
  instance Arity n => Vector (ContVec n) a where
    construct = accum
      (\(T_mkN f) a -> T_mkN (f . consPeano a))
      (\(T_mkN f)   -> toContVec $ f (CVecPeano unFun))
      (T_mkN id)
    inspect (ContVec c) f = c f
    {-# INLINE construct #-}
    {-# INLINE inspect   #-}
  
  newtype T_mkN n_tot a n = T_mkN (CVecPeano n a -> CVecPeano n_tot a)
  
  toContVec :: CVecPeano (Peano n) a -> ContVec n a
  toContVec = coerce
  
  newtype CVecPeano n a = CVecPeano (forall r. Fun n a r -> r)
  
  consPeano :: a -> CVecPeano n a -> CVecPeano ('S n) a
  consPeano a (CVecPeano cont) = CVecPeano $ \f -> cont $ curryFirst f a
  {-# INLINE consPeano #-}
  
  curryFirst :: Fun ('S n) a b -> a -> Fun n a b
  curryFirst = coerce
  {-# INLINE curryFirst #-}
  
  apply :: Arity n
        => (forall k. t ('S k) -> (a, t k))
        -> t (Peano n)
        -> ContVec n a
  {-# INLINE apply #-}
  apply step' z = toContVec $ fst (applyFun step' z)
  
  -----
  -- Data.Vector.Fixed.Mutable
  -----
  
  type family Mutable (v :: Type -> Type) :: Type -> Type -> Type
  
  type family DimM (v :: Type -> Type -> Type) :: Nat
  
  class (Arity (DimM v)) => MVector v a where
    new   :: PrimMonad m => m (v (PrimState m) a)
  
    unsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()
  
  class (Dim v ~ DimM (Mutable v), MVector (Mutable v) a) => IVector v a where
    unsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
  
    unsafeIndex :: v a -> Int -> a
  
  inspectVec :: forall v a b. (Arity (Dim v), IVector v a) => v a -> Fun (Peano (Dim v)) a b -> b
  {-# INLINE inspectVec #-}
  inspectVec v
    = inspect cv
    where
      cv :: ContVec (Dim v) a
      cv = apply (\(Const i) -> (unsafeIndex v i, Const (i+1)))
                 (Const 0 :: Const Int (Peano (Dim v)))
  
  constructVec :: forall v a. (Arity (Dim v), IVector v a) => Fun (Peano (Dim v)) a (v a)
  {-# INLINE constructVec #-}
  constructVec =
    accum step
          (\(T_new _ st) -> runST $ unsafeFreeze =<< st :: v a)
          (T_new 0 new :: T_new v a (Peano (Dim v)))
  
  data T_new v a n = T_new Int (forall s. ST s (Mutable v s a))
  
  step :: (IVector v a) => T_new v a ('S n) -> a -> T_new v a n
  step (T_new i st) x = T_new (i+1) $ do
    mv <- st
    unsafeWrite mv i x
    return mv
  
  -----
  -- Control.Monad.Primitive
  -----
  
  class Monad m => PrimMonad m where
    type PrimState m
  
  instance PrimMonad (ST s) where
    type PrimState (ST s) = s
  