{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Safe                  #-}
{-# LANGUAGE TypeFamilies          #-}
-- The following warning is disabled due to a necessary instance of Projectable
-- defined in this module.
{-# OPTIONS_GHC -fno-warn-orphans #-}

-- | Combinators to deal with streams carrying arrays.
module Copilot.Language.Operators.Array
  ( (!)
  , (!!)
  , (=:)
  , (=$)
  ) where

import Copilot.Core                          (Array, Op2 (Index),
                                              Op3 (UpdateArray), Typed, typeOf)
import Copilot.Language.Operators.Projection (Projectable(..))
import Copilot.Language.Stream               (Stream (..))

import Data.Word    (Word32)
import GHC.TypeLits (KnownNat)
import Prelude      hiding ((!!))

-- | Create a stream that carries an element of an array in another stream.
--
-- This function implements a projection of the element of an array at a given
-- position, over time. For example, if @s@ is a stream of type @Stream (Array
-- '5 Word8)@, then @s ! 3@ has type @Stream Word8@ and contains the 3rd
-- element (starting from zero) of the arrays in @s@ at any point in time.
(!) :: (KnownNat n, Typed t)
    => Stream (Array n t) -> Stream Word32 -> Stream t
Stream (Array n t)
arr ! :: forall (n :: Nat) t.
(KnownNat n, Typed t) =>
Stream (Array n t) -> Stream Word32 -> Stream t
! Stream Word32
n = Op2 (Array n t) Word32 t
-> Stream (Array n t) -> Stream Word32 -> Stream t
forall a1 b a.
(Typed a1, Typed b, Typed a) =>
Op2 a1 b a -> Stream a1 -> Stream b -> Stream a
Op2 (Type (Array n t) -> Op2 (Array n t) Word32 t
forall (n :: Nat) c. Type (Array n c) -> Op2 (Array n c) Word32 c
Index Type (Array n t)
forall a. Typed a => Type a
typeOf) Stream (Array n t)
arr Stream Word32
n

-- | Pair a stream with an element accessor, without applying it to obtain the
-- value of the element.
--
-- This function is needed to refer to an element accessor when the goal is to
-- update the element value, not just to read it.
(!!) :: Stream (Array n t)
     -> Stream Word32
     -> Projection (Array n t) (Stream Word32) t
!! :: forall (n :: Nat) t.
Stream (Array n t)
-> Stream Word32 -> Projection (Array n t) (Stream Word32) t
(!!) = Stream (Array n t)
-> Stream Word32 -> Projection (Array n t) (Stream Word32) t
forall (n :: Nat) t.
Stream (Array n t)
-> Stream Word32 -> Projection (Array n t) (Stream Word32) t
ProjectionA

-- | Update a stream of arrays.

-- This is an orphan instance; we suppress the warning that GHC would
-- normally produce with a GHC option at the top.
instance (KnownNat n, Typed t) => Projectable (Array n t) (Stream Word32) t where

  -- | A projection of an element of a stream of arrays.
  data Projection (Array n t) (Stream Word32) t =
       ProjectionA (Stream (Array n t)) (Stream Word32)

  -- | Create a stream where an element of an array has been updated with
  -- values from another stream.
  --
  -- For example, if an array has two elements of type @Int32@, and @s@ is a
  -- stream of such array type (@Stream (Array 2 Int32)@), and $v0$ is a stream
  -- of type @Int32@, then @s !! 0 =: v0@ has type @Stream (Array 2 Int32)@ and
  -- contains arrays where the value of the first element of each array is that
  -- of @v0@ at each point in time, and the value of the second element in the
  -- array is the same it had in @s@.
  =: :: Projection (Array n t) (Stream Word32) t
-> Stream t -> Stream (Array n t)
(=:) (ProjectionA Stream (Array n t)
s Stream Word32
ix) Stream t
v = Op3 (Array n t) Word32 t (Array n t)
-> Stream (Array n t)
-> Stream Word32
-> Stream t
-> Stream (Array n t)
forall a1 b c a.
(Typed a1, Typed b, Typed c, Typed a) =>
Op3 a1 b c a -> Stream a1 -> Stream b -> Stream c -> Stream a
Op3 (Type (Array n t) -> Op3 (Array n t) Word32 t (Array n t)
forall (n :: Nat) c.
Type (Array n c) -> Op3 (Array n c) Word32 c (Array n c)
UpdateArray Type (Array n t)
forall a. Typed a => Type a
typeOf) Stream (Array n t)
s Stream Word32
ix Stream t
v

  -- | Create a stream where an element of an array has been updated by
  -- applying a stream function to it.
  --
  -- For example, if an array has two elements of type @Int32@, and @s@ is a
  -- stream of such array type (@Stream (Array 2 Int32)@), and $f$ is function
  -- of type @Stream Int32 -> Stream Int32@, then @s !! 0 =$ f@ has type
  -- @Stream (Array 2 Int32)@ and contains arrays where the value of the first
  -- element of each array is that of @f (s ! 0)@ at each point in time, and
  -- the value of the second element in the array is the same it had in @s@.
  =$ :: Projection (Array n t) (Stream Word32) t
-> (Stream t -> Stream t) -> Stream (Array n t)
(=$) (ProjectionA Stream (Array n t)
s Stream Word32
ix) Stream t -> Stream t
op = Op3 (Array n t) Word32 t (Array n t)
-> Stream (Array n t)
-> Stream Word32
-> Stream t
-> Stream (Array n t)
forall a1 b c a.
(Typed a1, Typed b, Typed c, Typed a) =>
Op3 a1 b c a -> Stream a1 -> Stream b -> Stream c -> Stream a
Op3 (Type (Array n t) -> Op3 (Array n t) Word32 t (Array n t)
forall (n :: Nat) c.
Type (Array n c) -> Op3 (Array n c) Word32 c (Array n c)
UpdateArray Type (Array n t)
forall a. Typed a => Type a
typeOf) Stream (Array n t)
s Stream Word32
ix (Stream t -> Stream t
op (Stream (Array n t)
s Stream (Array n t) -> Stream Word32 -> Stream t
forall (n :: Nat) t.
(KnownNat n, Typed t) =>
Stream (Array n t) -> Stream Word32 -> Stream t
! Stream Word32
ix))