{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | Harpie arrays as 'Circuit.Stream' instances.
--
-- A non-empty array of shape @[n, m₁, …, mₖ]@ is a stream of @n@ rows, each
-- a token of shape @[m₁, …, mₖ]@.  The stream classes give the abstract
-- boundary-aware interface; the underlying operations are harpie's
-- 'Harpie.Array.cons', 'uncons', 'snoc', and 'unsnoc'.
--
-- The empty stream is 'Harpie.Array.empty'.  Because the shape of rows is
-- value-level in harpie, 'cons' and 'snoc' construct a one-row array when
-- applied to the empty stream.
module Circuit.Mat.Array.Stream
  ( -- * Re-exported classes
    Stream.Uncons (..),
    Stream.Cons (..),
    Stream.Snoc (..),
    Stream.These (..),
  )
where

import Circuit.Stream (Cons (..), Snoc (..), These (..), Uncons (..))
import Circuit.Stream qualified as Stream
import Data.Vector qualified as V
import Data.Vector.Unboxed qualified as VU
import Harpie.Array (Array (..))
import Harpie.Array qualified as D
import Prelude hiding (id, (.))

-- | An array is a stream of its rows.
--
-- * Empty array  → 'That' empty
-- * One-row array  → 'This' row
-- * Many-row array → 'These' head tail
instance Uncons (Array a) (Array a) where
  uncons :: Array a -> These (Array a) (Array a)
uncons Array a
a
    | Array a -> Bool
forall a. Array a -> Bool
D.isNull Array a
a = Array a -> These (Array a) (Array a)
forall a b. b -> These a b
That Array a
a
    | Array a -> Int
forall a. Array a -> Int
D.length Array a
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Array a -> These (Array a) (Array a)
forall a b. a -> These a b
This (Int -> Int -> Array a -> Array a
forall a. Int -> Int -> Array a -> Array a
D.select Int
0 Int
0 Array a
a)
    | Bool
otherwise =
        let (Array a
x, Array a
xs) = Array a -> (Array a, Array a)
forall a. Array a -> (Array a, Array a)
D.uncons Array a
a
         in Array a -> Array a -> These (Array a) (Array a)
forall a b. a -> b -> These a b
These Array a
x Array a
xs
  nil :: Array a
nil = Array a
forall a. Array a
D.empty

-- | Prepend a row.  When the tail is the empty stream, construct a
-- one-row array from the row's shape.
instance Cons (Array a) (Array a) where
  cons :: Array a -> Array a -> Array a
cons Array a
x Array a
xs
    | Array a -> Bool
forall a. Array a -> Bool
D.isNull Array a
xs = [Int] -> [a] -> Array a
forall t a. FromVector t a => [Int] -> t -> Array a
D.array (Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array a -> Vector Int
forall a. Array a -> Vector Int
D.shape Array a
x)) (Vector a -> [a]
forall a. Vector a -> [a]
V.toList (Array a -> Vector a
forall t a. FromVector t a => t -> Vector a
D.asVector Array a
x))
    | Bool
otherwise = Array a -> Array a -> Array a
forall a. Array a -> Array a -> Array a
D.cons Array a
x Array a
xs
  consNil :: Array a
consNil = Array a
forall a. Array a
D.empty

-- | Append a row.  When the initial stream is empty, construct a
-- one-row array from the row's shape.
instance Snoc (Array a) (Array a) where
  snoc :: Array a -> Array a -> Array a
snoc Array a
xs Array a
x
    | Array a -> Bool
forall a. Array a -> Bool
D.isNull Array a
xs = [Int] -> [a] -> Array a
forall t a. FromVector t a => [Int] -> t -> Array a
D.array (Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array a -> Vector Int
forall a. Array a -> Vector Int
D.shape Array a
x)) (Vector a -> [a]
forall a. Vector a -> [a]
V.toList (Array a -> Vector a
forall t a. FromVector t a => t -> Vector a
D.asVector Array a
x))
    | Bool
otherwise = Array a -> Array a -> Array a
forall a. Array a -> Array a -> Array a
D.snoc Array a
xs Array a
x
  snocNil :: Array a
snocNil = Array a
forall a. Array a
D.empty