{-# 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 :: forall a. Replicator a %1 -> ()
consume (Moved a
_) = ()
consume (Streamed ReplicationStream a
stream) = ReplicationStream a %1 -> ()
forall a. ReplicationStream a %1 -> ()
ReplicationStream.consume ReplicationStream a
stream
{-# INLINEABLE consume #-}

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

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

pure :: a -> Replicator a
pure :: forall a. a -> Replicator a
pure = a -> Replicator a
forall a. a -> Replicator a
Moved

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

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

liftA2 :: (a %1 -> b %1 -> c) -> Replicator a %1 -> Replicator b %1 -> Replicator c
liftA2 :: forall a b c.
(a %1 -> b %1 -> c)
-> Replicator a %1 -> Replicator b %1 -> Replicator c
liftA2 a %1 -> b %1 -> c
f (Moved a
a) (Moved b
b) = c -> Replicator c
forall a. a -> Replicator a
Moved (a %1 -> b %1 -> c
f a
a b
b)
liftA2 a %1 -> b %1 -> c
f (Moved a
a) (Streamed ReplicationStream b
s) = ReplicationStream c %1 -> Replicator c
forall a. ReplicationStream a -> Replicator a
Streamed ((b %1 -> c) -> ReplicationStream b %1 -> ReplicationStream c
forall a b.
(a %1 -> b) -> ReplicationStream a %1 -> ReplicationStream b
ReplicationStream.map (a %1 -> b %1 -> c
f a
a) ReplicationStream b
s)
liftA2 a %1 -> b %1 -> c
f (Streamed ReplicationStream a
s) (Moved b
b) = ReplicationStream c %1 -> Replicator c
forall a. ReplicationStream a -> Replicator a
Streamed ((a %1 -> c) -> ReplicationStream a %1 -> ReplicationStream c
forall a b.
(a %1 -> b) -> ReplicationStream a %1 -> ReplicationStream b
ReplicationStream.map (\a
a -> a %1 -> b %1 -> c
f a
a b
b) ReplicationStream a
s)
liftA2 a %1 -> b %1 -> c
f (Streamed ReplicationStream a
sa) (Streamed ReplicationStream b
sb) = ReplicationStream c %1 -> Replicator c
forall a. ReplicationStream a -> Replicator a
Streamed ((a %1 -> b %1 -> c)
-> ReplicationStream a
%1 -> ReplicationStream b
%1 -> ReplicationStream c
forall a b c.
(a %1 -> b %1 -> c)
-> ReplicationStream a
%1 -> ReplicationStream b
%1 -> ReplicationStream c
ReplicationStream.liftA2 a %1 -> b %1 -> c
f ReplicationStream a
sa ReplicationStream b
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 :: forall a. Replicator a %1 -> (a, Replicator a)
next (Moved a
x) = (a
x, a -> Replicator a
forall a. a -> Replicator a
Moved a
x)
next (Streamed (ReplicationStream s
s s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
consumes)) =
  s %1 -> (s, s)
dups s
s (s, s) %1 -> ((s, s) %1 -> (a, Replicator a)) -> (a, Replicator a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (s
s1, s
s2) -> (s %1 -> a
give s
s1, ReplicationStream a %1 -> Replicator a
forall a. ReplicationStream a -> Replicator a
Streamed (s
%1 -> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream s
s2 s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
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# :: forall a. Replicator a %1 -> (# a, Replicator a #)
next# (Moved a
x) = (# a
x, a -> Replicator a
forall a. a -> Replicator a
Moved a
x #)
next# (Streamed (ReplicationStream s
s s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
consumes)) =
  s %1 -> (s, s)
dups s
s (s, s)
%1 -> ((s, s) %1 -> (# a, Replicator a #)) -> (# a, Replicator a #)
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (s
s1, s
s2) -> (# s %1 -> a
give s
s1, ReplicationStream a %1 -> Replicator a
forall a. ReplicationStream a -> Replicator a
Streamed (s
%1 -> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream s
s2 s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
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 :: forall a. Int -> Replicator a %1 -> [a]
take Int
0 Replicator a
r =
  Replicator a %1 -> ()
forall a. Replicator a %1 -> ()
consume Replicator a
r () %1 -> (() %1 -> [a]) -> [a]
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    () -> []
take Int
1 Replicator a
r = [Replicator a %1 -> a
forall a. Replicator a %1 -> a
extract Replicator a
r]
take Int
n Replicator a
r =
  Replicator a %1 -> (a, Replicator a)
forall a. Replicator a %1 -> (a, Replicator a)
next Replicator a
r (a, Replicator a) %1 -> ((a, Replicator a) %1 -> [a]) -> [a]
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (a
a, Replicator a
r') -> a
a a %1 -> [a] %1 -> [a]
forall a. a -> [a] -> [a]
: Int -> Replicator a %1 -> [a]
forall a. Int -> Replicator a %1 -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Replicator a
r'

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

-- | Comonadic 'extend' function.
--
-- > extend f = map f . duplicate
extend :: (Replicator a %1 -> b) -> Replicator a %1 -> Replicator b
extend :: forall a b.
(Replicator a %1 -> b) -> Replicator a %1 -> Replicator b
extend Replicator a %1 -> b
f = (Replicator a %1 -> b)
-> Replicator (Replicator a) %1 -> Replicator b
forall a b. (a %1 -> b) -> Replicator a %1 -> Replicator b
map Replicator a %1 -> b
f (Replicator (Replicator a) %1 -> Replicator b)
-> (Replicator a %1 -> Replicator (Replicator a))
-> Replicator a
%1 -> Replicator b
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. Replicator a %1 -> Replicator (Replicator a)
forall a. Replicator a %1 -> Replicator (Replicator a)
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 :: 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
f Replicator a
r = forall (n :: Peano) a b.
Elim n a b =>
FunN n a b %1 -> Replicator a %1 -> b
elim' @(NatToPeano n) f
FunN (NatToPeano n) a b
f Replicator a
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' :: FunN 'Z a b %1 -> Replicator a %1 -> b
elim' FunN 'Z a b
b Replicator a
r =
    Replicator a %1 -> ()
forall a. Replicator a %1 -> ()
consume Replicator a
r () %1 -> (() %1 -> b) %1 -> b
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
      () -> b
FunN 'Z a b
b
  {-# INLINE elim' #-}

instance Elim ('S 'Z) a b where
  elim' :: FunN ('S 'Z) a b %1 -> Replicator a %1 -> b
elim' FunN ('S 'Z) a b
f Replicator a
r = FunN ('S 'Z) a b
a %1 -> b
f (Replicator a %1 -> a
forall a. Replicator a %1 -> a
extract Replicator a
r)
  {-# INLINE elim' #-}

instance (Elim ('S n) a b) => Elim ('S ('S n)) a b where
  elim' :: FunN ('S ('S n)) a b %1 -> Replicator a %1 -> b
elim' FunN ('S ('S n)) a b
g Replicator a
r =
    Replicator a %1 -> (a, Replicator a)
forall a. Replicator a %1 -> (a, Replicator a)
next Replicator a
r (a, Replicator a) %1 -> ((a, Replicator a) %1 -> b) %1 -> b
forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
      (a
a, Replicator a
r') -> forall (n :: Peano) a b.
Elim n a b =>
FunN n a b %1 -> Replicator a %1 -> b
elim' @('S n) (FunN ('S ('S n)) a b
a %1 -> a %1 -> FunN n a b
g a
a) Replicator a
r'
  {-# INLINE elim' #-}