# LANGUAGE AllowAmbiguousTypes #
# LANGUAGE DataKinds #
# LANGUAGE FlexibleContexts #
# LANGUAGE FlexibleInstances #
{-# LANGUAGE GADTs #-}
# LANGUAGE LambdaCase #
{-# LANGUAGE LinearTypes #-}
# LANGUAGE MagicHash #
# LANGUAGE MultiParamTypeClasses #
# LANGUAGE ScopedTypeVariables #
# LANGUAGE StandaloneKindSignatures #
# LANGUAGE TypeApplications #
# LANGUAGE TypeFamilies #
# LANGUAGE TypeOperators #
# LANGUAGE UnboxedTuples #
# LANGUAGE UndecidableInstances #
# LANGUAGE NoImplicitPrelude #
{-# OPTIONS_HADDOCK hide #-}

module Data.Replicator.Linear.Internal
  ( Replicator (..),
    consume,
    duplicate,
    map,
    pure,
    (<*>),
    liftA2,
    next,
    next#,
    take,
    extract,
    extend,
    Elim,
    elim,
  )
where

import Data.Arity.Linear.Internal
import Data.Kind (Constraint, Type)
import Data.Replicator.Linear.Internal.ReplicationStream (ReplicationStream (..))
import qualified Data.Replicator.Linear.Internal.ReplicationStream as ReplicationStream
import GHC.TypeLits
import Prelude.Linear.Internal
import Prelude ((-))
import qualified Prelude

-- | 'Replicator' is a stream-like data structure used to linearly duplicate
-- values.
data Replicator a where
  Moved :: a -> Replicator a
  Streamed :: ReplicationStream a %1 -> Replicator a

consume :: Replicator a %1 -> ()
consume (Moved _) = ()
consume (Streamed stream) = ReplicationStream.consume stream
# INLINEABLE consume #

duplicate :: Replicator a %1 -> Replicator (Replicator a)
duplicate = \case
  Moved x -> Moved (Moved x)
  Streamed stream -> Streamed $ ReplicationStream.map Streamed (ReplicationStream.duplicate stream)

map :: (a %1 -> b) -> Replicator a %1 -> Replicator b
map f = \case
  Moved x -> Moved (f x)
  Streamed stream -> Streamed $ ReplicationStream.map f stream

pure :: a -> Replicator a
pure = Moved

(<*>) :: Replicator (a %1 -> b) %1 -> Replicator a %1 -> Replicator b
Moved f <*> Moved x = Moved (f x)
Moved f <*> Streamed s = Streamed (ReplicationStream.map f s)
Streamed fs <*> Moved x = Streamed (ReplicationStream.map (\f -> f x) fs)
Streamed sf <*> Streamed sx = Streamed (sf ReplicationStream.<*> sx)

infixl 4 <*> -- same fixity as base.<*>

liftA2 :: (a %1 -> b %1 -> c) -> Replicator a %1 -> Replicator b %1 -> Replicator c
liftA2 f (Moved a) (Moved b) = Moved (f a b)
liftA2 f (Moved a) (Streamed s) = Streamed (ReplicationStream.map (f a) s)
liftA2 f (Streamed s) (Moved b) = Streamed (ReplicationStream.map (\a -> f a b) s)
liftA2 f (Streamed sa) (Streamed sb) = Streamed (ReplicationStream.liftA2 f sa sb)
-- We need to inline this to get good results with generic deriving of
  Dupable .
# INLINE liftA2 #

-- | Extracts the next item from the \"infinite stream\" @'Replicator' a@.
next :: Replicator a %1 -> (a, Replicator a)
next (Moved x) = (x, Moved x)
next (Streamed (ReplicationStream s give dups consumes)) =
  dups s & \case
    (s1, s2) -> (give s1, Streamed (ReplicationStream s2 give dups consumes))
# INLINEABLE next #

-- | Extracts the next item from the \"infinite stream\" @'Replicator' a@.
-- Same function as 'next', but returning an unboxed tuple.
next# :: Replicator a %1 -> (# a, Replicator a #)
next# (Moved x) = (# x, Moved x #)
next# (Streamed (ReplicationStream s give dups consumes)) =
  dups s & \case
    (s1, s2) -> (# give s1, Streamed (ReplicationStream s2 give dups consumes) #)
# INLINEABLE next # #

  | @'take ' n as@ is a list of size @n@ , containing @n@ replicas from @as@.
take :: Prelude.Int -> Replicator a %1 -> [a]
take 0 r =
  consume r & \case
    () -> []
take 1 r = [extract r]
take n r =
  next r & \case
    (a, r') -> a : take (n - 1) r'

-- | Returns the next item from @'Replicator' a@ and efficiently consumes
-- the replicator at the same time.
extract :: Replicator a %1 -> a
extract (Moved x) = x
extract (Streamed (ReplicationStream s give _ _)) = give s
# INLINEABLE extract #

-- | Comonadic 'extend' function.
--
-- > extend f = map f . duplicate
extend :: (Replicator a %1 -> b) -> Replicator a %1 -> Replicator b
extend f = map f . duplicate

  | Takes a function of type @a % 1 - > a % 1 - > ... % 1 - > a % 1 - > b@ , and
  returns a @b@ . The replicator is used to supply all the items of type @a@
-- required by the function.
--
-- For instance:
--
-- > elim @1 :: (a %1 -> b) %1 -> Replicator a %1 -> b
-- > elim @2 :: (a %1 -> a %1 -> b) %1 -> Replicator a %1 -> b
-- > elim @3 :: (a %1 -> a %1 -> a %1 -> b) %1 -> Replicator a %1 -> b
--
-- It is not always necessary to give the arity argument. It can be
-- inferred from the function argument.
--
-- > elim (,) :: Replicator a %1 -> (a, a)
-- > elim (,,) :: Replicator a %1 -> (a, a, a)
--
-- About the constraints of this function (they won't get in your way):
--
  * @'Elim ' ( ' NatToPeano ' n ) a b@ provides the actual implementation of ' elim ' ; there is an instance of this class for any @(n , a , b)@
  * @'IsFunN ' a b f , f ~ ' FunN ' ( ' NatToPeano ' n ) a b , n ~ ' Arity ' b f@ indicate the shape of @f@ to the typechecker ( see documentation of ' IsFunN ' ) .
elim ::
  forall (n :: Nat) a b f.
  ( Elim (NatToPeano n) a b,
    IsFunN a b f,
    f ~ FunN (NatToPeano n) a b,
    n ~ Arity b f
  ) =>
  f %1 ->
  Replicator a %1 ->
  b
elim f r = elim' @(NatToPeano n) f r

-- | @'Elim' n a b@ is used to implement 'elim' without recursion
-- so that we can guarantee that 'elim' will be inlined and unrolled.
--
-- 'Elim' is solely used in the signature of 'elim'.
type Elim :: Peano -> Type -> Type -> Constraint
class Elim n a b where
  -- Note that 'elim' is, in particular, used to force eta-expansion of
  -- 'elim\''.  Otherwise, 'elim\'' might not get inlined (see
  -base/issues/369 ) .
  elim' :: FunN n a b %1 -> Replicator a %1 -> b

instance Elim 'Z a b where
  elim' b r =
    consume r & \case
      () -> b
# INLINE elim ' #

instance Elim ('S 'Z) a b where
  elim' f r = f (extract r)
# INLINE elim ' #

instance (Elim ('S n) a b) => Elim ('S ('S n)) a b where
  elim' g r =
    next r & \case
      (a, r') -> elim' @('S n) (g a) r'
# INLINE elim ' #

# LANGUAGE GADTs #
# LANGUAGE LinearTypes #
# OPTIONS_HADDOCK hide #
 | 'Replicator' is a stream-like data structure used to linearly duplicate
 values.
 same fixity as base.<*>
 We need to inline this to get good results with generic deriving of
 | Extracts the next item from the \"infinite stream\" @'Replicator' a@.
 | Extracts the next item from the \"infinite stream\" @'Replicator' a@.
 Same function as 'next', but returning an unboxed tuple.
 | Returns the next item from @'Replicator' a@ and efficiently consumes
 the replicator at the same time.
 | Comonadic 'extend' function.

 > extend f = map f . duplicate
 required by the function.

 For instance:

 > elim @1 :: (a %1 -> b) %1 -> Replicator a %1 -> b
 > elim @2 :: (a %1 -> a %1 -> b) %1 -> Replicator a %1 -> b
 > elim @3 :: (a %1 -> a %1 -> a %1 -> b) %1 -> Replicator a %1 -> b

 It is not always necessary to give the arity argument. It can be
 inferred from the function argument.

 > elim (,) :: Replicator a %1 -> (a, a)
 > elim (,,) :: Replicator a %1 -> (a, a, a)

 About the constraints of this function (they won't get in your way):

 | @'Elim' n a b@ is used to implement 'elim' without recursion
 so that we can guarantee that 'elim' will be inlined and unrolled.

 'Elim' is solely used in the signature of 'elim'.
 Note that 'elim' is, in particular, used to force eta-expansion of
 'elim\''.  Otherwise, 'elim\'' might not get inlined (see
# LANGUAGE AllowAmbiguousTypes #
# LANGUAGE DataKinds #
# LANGUAGE FlexibleContexts #
# LANGUAGE FlexibleInstances #
# LANGUAGE LambdaCase #
# LANGUAGE MagicHash #
# LANGUAGE MultiParamTypeClasses #
# LANGUAGE ScopedTypeVariables #
# LANGUAGE StandaloneKindSignatures #
# LANGUAGE TypeApplications #
# LANGUAGE TypeFamilies #
# LANGUAGE TypeOperators #
# LANGUAGE UnboxedTuples #
# LANGUAGE UndecidableInstances #
# LANGUAGE NoImplicitPrelude #

module Data.Replicator.Linear.Internal
  ( Replicator (..),
    consume,
    duplicate,
    map,
    pure,
    (<*>),
    liftA2,
    next,
    next#,
    take,
    extract,
    extend,
    Elim,
    elim,
  )
where

import Data.Arity.Linear.Internal
import Data.Kind (Constraint, Type)
import Data.Replicator.Linear.Internal.ReplicationStream (ReplicationStream (..))
import qualified Data.Replicator.Linear.Internal.ReplicationStream as ReplicationStream
import GHC.TypeLits
import Prelude.Linear.Internal
import Prelude ((-))
import qualified Prelude

data Replicator a where
  Moved :: a -> Replicator a
  Streamed :: ReplicationStream a %1 -> Replicator a

consume :: Replicator a %1 -> ()
consume (Moved _) = ()
consume (Streamed stream) = ReplicationStream.consume stream
# INLINEABLE consume #

duplicate :: Replicator a %1 -> Replicator (Replicator a)
duplicate = \case
  Moved x -> Moved (Moved x)
  Streamed stream -> Streamed $ ReplicationStream.map Streamed (ReplicationStream.duplicate stream)

map :: (a %1 -> b) -> Replicator a %1 -> Replicator b
map f = \case
  Moved x -> Moved (f x)
  Streamed stream -> Streamed $ ReplicationStream.map f stream

pure :: a -> Replicator a
pure = Moved

(<*>) :: Replicator (a %1 -> b) %1 -> Replicator a %1 -> Replicator b
Moved f <*> Moved x = Moved (f x)
Moved f <*> Streamed s = Streamed (ReplicationStream.map f s)
Streamed fs <*> Moved x = Streamed (ReplicationStream.map (\f -> f x) fs)
Streamed sf <*> Streamed sx = Streamed (sf ReplicationStream.<*> sx)


liftA2 :: (a %1 -> b %1 -> c) -> Replicator a %1 -> Replicator b %1 -> Replicator c
liftA2 f (Moved a) (Moved b) = Moved (f a b)
liftA2 f (Moved a) (Streamed s) = Streamed (ReplicationStream.map (f a) s)
liftA2 f (Streamed s) (Moved b) = Streamed (ReplicationStream.map (\a -> f a b) s)
liftA2 f (Streamed sa) (Streamed sb) = Streamed (ReplicationStream.liftA2 f sa sb)
  Dupable .
# INLINE liftA2 #

next :: Replicator a %1 -> (a, Replicator a)
next (Moved x) = (x, Moved x)
next (Streamed (ReplicationStream s give dups consumes)) =
  dups s & \case
    (s1, s2) -> (give s1, Streamed (ReplicationStream s2 give dups consumes))
# INLINEABLE next #

next# :: Replicator a %1 -> (# a, Replicator a #)
next# (Moved x) = (# x, Moved x #)
next# (Streamed (ReplicationStream s give dups consumes)) =
  dups s & \case
    (s1, s2) -> (# give s1, Streamed (ReplicationStream s2 give dups consumes) #)
# INLINEABLE next # #

  | @'take ' n as@ is a list of size @n@ , containing @n@ replicas from @as@.
take :: Prelude.Int -> Replicator a %1 -> [a]
take 0 r =
  consume r & \case
    () -> []
take 1 r = [extract r]
take n r =
  next r & \case
    (a, r') -> a : take (n - 1) r'

extract :: Replicator a %1 -> a
extract (Moved x) = x
extract (Streamed (ReplicationStream s give _ _)) = give s
# INLINEABLE extract #

extend :: (Replicator a %1 -> b) -> Replicator a %1 -> Replicator b
extend f = map f . duplicate

  | Takes a function of type @a % 1 - > a % 1 - > ... % 1 - > a % 1 - > b@ , and
  returns a @b@ . The replicator is used to supply all the items of type @a@
  * @'Elim ' ( ' NatToPeano ' n ) a b@ provides the actual implementation of ' elim ' ; there is an instance of this class for any @(n , a , b)@
  * @'IsFunN ' a b f , f ~ ' FunN ' ( ' NatToPeano ' n ) a b , n ~ ' Arity ' b f@ indicate the shape of @f@ to the typechecker ( see documentation of ' IsFunN ' ) .
elim ::
  forall (n :: Nat) a b f.
  ( Elim (NatToPeano n) a b,
    IsFunN a b f,
    f ~ FunN (NatToPeano n) a b,
    n ~ Arity b f
  ) =>
  f %1 ->
  Replicator a %1 ->
  b
elim f r = elim' @(NatToPeano n) f r

type Elim :: Peano -> Type -> Type -> Constraint
class Elim n a b where
  -base/issues/369 ) .
  elim' :: FunN n a b %1 -> Replicator a %1 -> b

instance Elim 'Z a b where
  elim' b r =
    consume r & \case
      () -> b
# INLINE elim ' #

instance Elim ('S 'Z) a b where
  elim' f r = f (extract r)
# INLINE elim ' #

instance (Elim ('S n) a b) => Elim ('S ('S n)) a b where
  elim' g r =
    next r & \case
      (a, r') -> elim' @('S n) (g a) r'
# INLINE elim ' #
