content: {-# 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)) =
    case dups s of
      (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)) =
    case dups s of
      (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 =
    case consume r of
      () -> []
  take 1 r = [extract r]
  take n r =
    case next r of
      (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
    -- https://github.com/tweag/linear-base/issues/369).
    elim' :: FunN n a b %1 -> Replicator a %1 -> b
  
  instance Elim 'Z a b where
    elim' b r =
      case consume r of
        () -> 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 =
      case next r of
        (a, r') -> elim' @('S n) (g a) r'
    {-# INLINE elim' #-}
  