{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Test.StateMachine.Lockstep.Auxiliary (
    Elem(..)
  , npAt
  , NTraversable(..)
  , ntraverse
  , ncfmap
  , nfmap
  , ncfoldMap
  , nfoldMap
  ) where

import           Control.Monad.State
import           Data.Kind
                   (Type)
import           Prelude
#if !MIN_VERSION_base(4,13,0)
import           Data.Monoid
                   (Monoid)
#endif
import           Data.Proxy
import           Data.SOP

{-------------------------------------------------------------------------------
  Auxiliary
-------------------------------------------------------------------------------}

-- TODO: Could simulate by using @NS xs (K ())@
data Elem (xs :: [k]) (a :: k) where
  ElemHead :: Elem (k ': ks) k
  ElemTail :: Elem ks k -> Elem (k' ': ks) k

npAt' :: Elem xs a -> NP f xs -> f a
npAt' :: Elem xs a -> NP f xs -> f a
npAt' Elem xs a
ElemHead      (f x
f :* NP f xs
_)  = f a
f x
f
npAt' (ElemTail Elem ks a
ix) (f x
_ :* NP f xs
fs) = Elem ks a -> NP f ks -> f a
forall k (xs :: [k]) (a :: k) (f :: k -> *).
Elem xs a -> NP f xs -> f a
npAt' Elem ks a
ix NP f ks
NP f xs
fs

npAt :: NP f xs -> Elem xs a -> f a
npAt :: NP f xs -> Elem xs a -> f a
npAt = (Elem xs a -> NP f xs -> f a) -> NP f xs -> Elem xs a -> f a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Elem xs a -> NP f xs -> f a
forall k (xs :: [k]) (a :: k) (f :: k -> *).
Elem xs a -> NP f xs -> f a
npAt'

-- | N-ary traversable functors
--
-- TODO: Don't provide Elem explicitly (just instantiate @c@)?
-- TODO: Introduce HTraverse into SOP?
class NTraversable (f :: (k -> Type) -> [k] -> Type) where
  nctraverse :: (Applicative m, All c xs)
             => proxy c
             -> (forall a. c a => Elem xs a -> g a -> m (h a))
             -> f g xs -> m (f h xs)

ntraverse :: (NTraversable f, Applicative m, SListI xs)
          => (forall a. Elem xs a -> g a -> m (h a))
          -> f g xs -> m (f h xs)
ntraverse :: (forall (a :: k). Elem xs a -> g a -> m (h a))
-> f g xs -> m (f h xs)
ntraverse = Proxy Top
-> (forall (a :: k). Top a => Elem xs a -> g a -> m (h a))
-> f g xs
-> m (f h xs)
forall k (f :: (k -> *) -> [k] -> *) (m :: * -> *)
       (c :: k -> Constraint) (xs :: [k])
       (proxy :: (k -> Constraint) -> *) (g :: k -> *) (h :: k -> *).
(NTraversable f, Applicative m, All c xs) =>
proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> m (h a))
-> f g xs
-> m (f h xs)
nctraverse (Proxy Top
forall k (t :: k). Proxy t
Proxy @Top)

ncfmap :: (NTraversable f, All c xs)
       => proxy c
       -> (forall a. c a => Elem xs a -> g a -> h a)
       -> f g xs -> f h xs
ncfmap :: proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> h a)
-> f g xs
-> f h xs
ncfmap proxy c
p forall (a :: k). c a => Elem xs a -> g a -> h a
f f g xs
xs = I (f h xs) -> f h xs
forall a. I a -> a
unI (I (f h xs) -> f h xs) -> I (f h xs) -> f h xs
forall a b. (a -> b) -> a -> b
$ proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> I (h a))
-> f g xs
-> I (f h xs)
forall k (f :: (k -> *) -> [k] -> *) (m :: * -> *)
       (c :: k -> Constraint) (xs :: [k])
       (proxy :: (k -> Constraint) -> *) (g :: k -> *) (h :: k -> *).
(NTraversable f, Applicative m, All c xs) =>
proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> m (h a))
-> f g xs
-> m (f h xs)
nctraverse proxy c
p (\Elem xs a
ix -> h a -> I (h a)
forall a. a -> I a
I (h a -> I (h a)) -> (g a -> h a) -> g a -> I (h a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Elem xs a -> g a -> h a
forall (a :: k). c a => Elem xs a -> g a -> h a
f Elem xs a
ix) f g xs
xs

nfmap :: (NTraversable f, SListI xs)
      => (forall a. Elem xs a -> g a -> h a)
      -> f g xs -> f h xs
nfmap :: (forall (a :: k). Elem xs a -> g a -> h a) -> f g xs -> f h xs
nfmap forall (a :: k). Elem xs a -> g a -> h a
f f g xs
xs = Proxy Top
-> (forall (a :: k). Top a => Elem xs a -> g a -> h a)
-> f g xs
-> f h xs
forall k (f :: (k -> *) -> [k] -> *) (c :: k -> Constraint)
       (xs :: [k]) (proxy :: (k -> Constraint) -> *) (g :: k -> *)
       (h :: k -> *).
(NTraversable f, All c xs) =>
proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> h a)
-> f g xs
-> f h xs
ncfmap (Proxy Top
forall k (t :: k). Proxy t
Proxy @Top) forall (a :: k). Top a => Elem xs a -> g a -> h a
forall (a :: k). Elem xs a -> g a -> h a
f f g xs
xs

ncfoldMap :: forall proxy f g m c xs.
             (NTraversable f, Monoid m, All c xs)
          => proxy c
          -> (forall a. c a => Elem xs a -> g a -> m)
          -> f g xs -> m
ncfoldMap :: proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> m) -> f g xs -> m
ncfoldMap proxy c
p forall (a :: k). c a => Elem xs a -> g a -> m
f = \f g xs
xs -> State m (f g xs) -> m -> m
forall s a. State s a -> s -> s
execState (f g xs -> State m (f g xs)
aux f g xs
xs) m
forall a. Monoid a => a
mempty
  where
    aux :: f g xs -> State m (f g xs)
    aux :: f g xs -> State m (f g xs)
aux f g xs
xs = proxy c
-> (forall (a :: k).
    c a =>
    Elem xs a -> g a -> StateT m Identity (g a))
-> f g xs
-> State m (f g xs)
forall k (f :: (k -> *) -> [k] -> *) (m :: * -> *)
       (c :: k -> Constraint) (xs :: [k])
       (proxy :: (k -> Constraint) -> *) (g :: k -> *) (h :: k -> *).
(NTraversable f, Applicative m, All c xs) =>
proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> m (h a))
-> f g xs
-> m (f h xs)
nctraverse proxy c
p forall (a :: k). c a => Elem xs a -> g a -> StateT m Identity (g a)
aux' f g xs
xs

    aux' :: c a => Elem xs a -> g a -> State m (g a)
    aux' :: Elem xs a -> g a -> State m (g a)
aux' Elem xs a
ix g a
ga = (m -> m) -> StateT m Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Elem xs a -> g a -> m
forall (a :: k). c a => Elem xs a -> g a -> m
f Elem xs a
ix g a
ga m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend`) StateT m Identity () -> State m (g a) -> State m (g a)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> g a -> State m (g a)
forall (m :: * -> *) a. Monad m => a -> m a
return g a
ga

nfoldMap :: (NTraversable f, Monoid m, SListI xs)
         => (forall a. Elem xs a -> g a -> m)
         -> f g xs -> m
nfoldMap :: (forall (a :: k). Elem xs a -> g a -> m) -> f g xs -> m
nfoldMap forall (a :: k). Elem xs a -> g a -> m
f f g xs
xs = Proxy Top
-> (forall (a :: k). Top a => Elem xs a -> g a -> m) -> f g xs -> m
forall k (proxy :: (k -> Constraint) -> *)
       (f :: (k -> *) -> [k] -> *) (g :: k -> *) m (c :: k -> Constraint)
       (xs :: [k]).
(NTraversable f, Monoid m, All c xs) =>
proxy c
-> (forall (a :: k). c a => Elem xs a -> g a -> m) -> f g xs -> m
ncfoldMap (Proxy Top
forall k (t :: k). Proxy t
Proxy @Top) forall (a :: k). Top a => Elem xs a -> g a -> m
forall (a :: k). Elem xs a -> g a -> m
f f g xs
xs