{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_HADDOCK hide #-}

module Data.Replicator.Linear.Internal.ReplicationStream
  ( ReplicationStream (..),
    consume,
    duplicate,
    map,
    pure,
    (<*>),
    liftA2,
  )
where

import Data.Unrestricted.Linear.Internal.Ur
import Prelude.Linear.Internal

-- | @ReplicationStream s g dup2 c@ is the infinite linear stream
-- @repeat (g s)@ where @dup2@ is used to make as many copies of @s@ as
-- necessary, and @c@ is used to consume @s@ when consuming the stream.
--
-- Although it isn't enforced at type level, @dup2@ should abide by the same
-- laws as 'Data.Unrestricted.Linear.dup2':
-- * @first c (dup2 a) ≃ a ≃ second c (dup2 a)@ (neutrality)
-- * @first dup2 (dup2 a) ≃ (second dup2 (dup2 a))@ (associativity)
--
-- This type is solely used to implement 'Data.Replicator.Linear'
data ReplicationStream a where
  ReplicationStream ::
    s %1 ->
    (s %1 -> a) ->
    (s %1 -> (s, s)) ->
    (s %1 -> ()) ->
    ReplicationStream a

consume :: ReplicationStream a %1 -> ()
consume :: forall a. ReplicationStream a %1 -> ()
consume (ReplicationStream s
s s %1 -> a
_ s %1 -> (s, s)
_ s %1 -> ()
consumes) = s %1 -> ()
consumes s
s
{-# INLINEABLE consume #-}

duplicate :: ReplicationStream a %1 -> ReplicationStream (ReplicationStream a)
duplicate :: forall a.
ReplicationStream a %1 -> ReplicationStream (ReplicationStream a)
duplicate (ReplicationStream s
s s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
consumes) =
  forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream
    s
s
    (\s
s' -> forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream s
s' s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
consumes)
    s %1 -> (s, s)
dups
    s %1 -> ()
consumes

map :: (a %1 -> b) -> ReplicationStream a %1 -> ReplicationStream b
map :: forall a b.
(a %1 -> b) -> ReplicationStream a %1 -> ReplicationStream b
map a %1 -> b
f (ReplicationStream s
s s %1 -> a
give s %1 -> (s, s)
dups s %1 -> ()
consumes) =
  forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream s
s (a %1 -> b
f forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. s %1 -> a
give) s %1 -> (s, s)
dups s %1 -> ()
consumes

pure :: a -> ReplicationStream a
pure :: forall a. a -> ReplicationStream a
pure a
x =
  forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream
    (forall a. a -> Ur a
Ur a
x)
    forall a. Ur a %1 -> a
unur
    ( \case
        Ur a
x' -> (forall a. a -> Ur a
Ur a
x', forall a. a -> Ur a
Ur a
x')
    )
    ( \case
        Ur a
_ -> ()
    )

(<*>) :: ReplicationStream (a %1 -> b) %1 -> ReplicationStream a %1 -> ReplicationStream b
(ReplicationStream s
sf s %1 -> a %1 -> b
givef s %1 -> (s, s)
dupsf s %1 -> ()
consumesf) <*> :: forall a b.
ReplicationStream (a %1 -> b)
%1 -> ReplicationStream a %1 -> ReplicationStream b
<*> (ReplicationStream s
sx s %1 -> a
givex s %1 -> (s, s)
dupsx s %1 -> ()
consumesx) =
  forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream
    (s
sf, s
sx)
    (\(s
sf', s
sx') -> s %1 -> a %1 -> b
givef s
sf' (s %1 -> a
givex s
sx'))
    ( \(s
sf', s
sx') ->
        (s %1 -> (s, s)
dupsf s
sf', s %1 -> (s, s)
dupsx s
sx') forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
          ((s
sf1, s
sf2), (s
sx1, s
sx2)) -> ((s
sf1, s
sx1), (s
sf2, s
sx2))
    )
    ( \(s
sf', s
sx') ->
        s %1 -> ()
consumesf s
sf' forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
          () -> s %1 -> ()
consumesx s
sx'
    )

liftA2 :: (a %1 -> b %1 -> c) -> ReplicationStream a %1 -> ReplicationStream b %1 -> ReplicationStream c
liftA2 :: forall a b c.
(a %1 -> b %1 -> c)
-> ReplicationStream a
%1 -> ReplicationStream b
%1 -> ReplicationStream c
liftA2 a %1 -> b %1 -> c
f (ReplicationStream s
sa s %1 -> a
givea s %1 -> (s, s)
dupsa s %1 -> ()
consumesa) (ReplicationStream s
sb s %1 -> b
giveb s %1 -> (s, s)
dupsb s %1 -> ()
consumesb) =
  forall s a.
s
-> (s %1 -> a)
-> (s %1 -> (s, s))
-> (s %1 -> ())
-> ReplicationStream a
ReplicationStream
    (s
sa, s
sb)
    (\(s
sa', s
sb') -> a %1 -> b %1 -> c
f (s %1 -> a
givea s
sa') (s %1 -> b
giveb s
sb'))
    ( \(s
sa', s
sb') ->
        (s %1 -> (s, s)
dupsa s
sa', s %1 -> (s, s)
dupsb s
sb') forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
          ((s
sa1, s
sa2), (s
sb1, s
sb2)) -> ((s
sa1, s
sb1), (s
sa2, s
sb2))
    )
    ( \(s
sa', s
sb') ->
        s %1 -> ()
consumesa s
sa' forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
          () -> s %1 -> ()
consumesb s
sb'
    )
-- We need to inline this to get good results with generic deriving
-- of Dupable.
{-# INLINE liftA2 #-}

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