module Data.RRBVector.Internal.Indexed where

-- TODO: use unboxed tuples?

data WithIndex a = WithIndex !Int a

-- | > Compose (State Int) f a
newtype Indexed f a = Indexed { Indexed f a -> Int -> WithIndex (f a)
runIndexed :: Int -> WithIndex (f a) }

instance Functor f => Functor (Indexed f) where
    fmap :: (a -> b) -> Indexed f a -> Indexed f b
fmap a -> b
f (Indexed Int -> WithIndex (f a)
sf) = (Int -> WithIndex (f b)) -> Indexed f b
forall (f :: * -> *) a. (Int -> WithIndex (f a)) -> Indexed f a
Indexed ((Int -> WithIndex (f b)) -> Indexed f b)
-> (Int -> WithIndex (f b)) -> Indexed f b
forall a b. (a -> b) -> a -> b
$ \Int
s -> let WithIndex Int
s' f a
x = Int -> WithIndex (f a)
sf Int
s in Int -> f b -> WithIndex (f b)
forall a. Int -> a -> WithIndex a
WithIndex Int
s' ((a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f f a
x)
    {-# INLINE fmap #-}

instance Applicative f => Applicative (Indexed f) where
    pure :: a -> Indexed f a
pure a
x = (Int -> WithIndex (f a)) -> Indexed f a
forall (f :: * -> *) a. (Int -> WithIndex (f a)) -> Indexed f a
Indexed ((Int -> WithIndex (f a)) -> Indexed f a)
-> (Int -> WithIndex (f a)) -> Indexed f a
forall a b. (a -> b) -> a -> b
$ \Int
s -> Int -> f a -> WithIndex (f a)
forall a. Int -> a -> WithIndex a
WithIndex Int
s (a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
    {-# INLINE pure #-}

    Indexed Int -> WithIndex (f (a -> b))
sfa <*> :: Indexed f (a -> b) -> Indexed f a -> Indexed f b
<*> Indexed Int -> WithIndex (f a)
sfb = (Int -> WithIndex (f b)) -> Indexed f b
forall (f :: * -> *) a. (Int -> WithIndex (f a)) -> Indexed f a
Indexed ((Int -> WithIndex (f b)) -> Indexed f b)
-> (Int -> WithIndex (f b)) -> Indexed f b
forall a b. (a -> b) -> a -> b
$ \Int
s ->
        let WithIndex Int
s' f (a -> b)
f = Int -> WithIndex (f (a -> b))
sfa Int
s
            WithIndex Int
s'' f a
x = Int -> WithIndex (f a)
sfb Int
s'
        in Int -> f b -> WithIndex (f b)
forall a. Int -> a -> WithIndex a
WithIndex Int
s'' (f (a -> b)
f f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f a
x)
    {-# INLINE (<*>) #-}

evalIndexed :: Indexed f a -> Int -> f a
evalIndexed :: Indexed f a -> Int -> f a
evalIndexed (Indexed Int -> WithIndex (f a)
sf) Int
x = let WithIndex Int
_ f a
y = Int -> WithIndex (f a)
sf Int
x in f a
y
{-# INLINE evalIndexed #-}