-- SPDX-License-Identifier: MPL-2.0

{- | Dense matrices with finite, row-major storage and zero-based indices.

The constructor is hidden so the stored element count always agrees with the
shape. Zero rows and zero columns are valid and retain their shape. Checked
operations describe invalid inputs with 'MathError'; the corresponding
'Maybe' operations retain the original failure convention.
-}
module HMath.Matrix (
    Matrix,
    fromList,
    fromListChecked,
    toList,
    shape,
    zeroMatrix,
    zeroMatrixChecked,
    diagonal,
    diagonalChecked,
    at,
    atChecked,
    sliceMatrix,
    sliceMatrixChecked,
    reshape,
    reshapeChecked,
    addMatrix,
    addMatrixChecked,
    subtractMatrix,
    subtractMatrixChecked,
    hadamard,
    hadamardChecked,
    scaleMatrix,
    multiply,
    multiplyChecked,
    transposeMatrix,
    identity,
    identityChecked,
) where

import Data.List (transpose)
import HMath.Error (MathError (..))

{- | An abstract dense matrix with explicit dimensions and flat row-major
storage. Empty shapes require no list of empty rows. 'Show' displays the
internal representation for inspection; it is not a serialization format.
-}
data Matrix a = Matrix Int Int [a]
    deriving (Matrix a -> Matrix a -> Bool
(Matrix a -> Matrix a -> Bool)
-> (Matrix a -> Matrix a -> Bool) -> Eq (Matrix a)
forall a. Eq a => Matrix a -> Matrix a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Matrix a -> Matrix a -> Bool
== :: Matrix a -> Matrix a -> Bool
$c/= :: forall a. Eq a => Matrix a -> Matrix a -> Bool
/= :: Matrix a -> Matrix a -> Bool
Eq, Int -> Matrix a -> ShowS
[Matrix a] -> ShowS
Matrix a -> String
(Int -> Matrix a -> ShowS)
-> (Matrix a -> String) -> ([Matrix a] -> ShowS) -> Show (Matrix a)
forall a. Show a => Int -> Matrix a -> ShowS
forall a. Show a => [Matrix a] -> ShowS
forall a. Show a => Matrix a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Matrix a -> ShowS
showsPrec :: Int -> Matrix a -> ShowS
$cshow :: forall a. Show a => Matrix a -> String
show :: Matrix a -> String
$cshowList :: forall a. Show a => [Matrix a] -> ShowS
showList :: [Matrix a] -> ShowS
Show)

{- | Construct a matrix from a finite list in row-major order. Negative shapes,
shapes whose element count exceeds @maxBound :: Int@, and incorrect list
lengths fail.
-}
fromList :: Int -> Int -> [a] -> Maybe (Matrix a)
fromList :: forall a. Int -> Int -> [a] -> Maybe (Matrix a)
fromList Int
rows Int
cols = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Either MathError (Matrix a) -> Maybe (Matrix a))
-> ([a] -> Either MathError (Matrix a)) -> [a] -> Maybe (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked Int
rows Int
cols

-- | Checked 'fromList'. The shape is validated before inspecting the list.
fromListChecked :: Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked :: forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
fromListChecked Int
rows Int
cols [a]
values = do
    expected <- Int -> Int -> Either MathError Int
elementCount Int
rows Int
cols
    let actual = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values
    if expected == actual
        then Right (Matrix rows cols values)
        else Left (ElementCountMismatch expected actual)

-- | Return all elements in row-major order.
toList :: Matrix a -> [a]
toList :: forall a. Matrix a -> [a]
toList (Matrix Int
_ Int
_ [a]
values) = [a]
values

-- | Return @(rows, columns)@, including either dimension of an empty matrix.
shape :: Matrix a -> (Int, Int)
shape :: forall a. Matrix a -> (Int, Int)
shape (Matrix Int
rows Int
cols [a]
_) = (Int
rows, Int
cols)

{- | Construct a matrix filled with zeros, subject to the same shape checks as
'fromList'.
-}
zeroMatrix :: (Num a) => Int -> Int -> Maybe (Matrix a)
zeroMatrix :: forall a. Num a => Int -> Int -> Maybe (Matrix a)
zeroMatrix Int
rows Int
cols = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Int -> Int -> Either MathError (Matrix a)
forall a. Num a => Int -> Int -> Either MathError (Matrix a)
zeroMatrixChecked Int
rows Int
cols)

-- | Checked 'zeroMatrix'.
zeroMatrixChecked :: (Num a) => Int -> Int -> Either MathError (Matrix a)
zeroMatrixChecked :: forall a. Num a => Int -> Int -> Either MathError (Matrix a)
zeroMatrixChecked Int
rows Int
cols = do
    count <- Int -> Int -> Either MathError Int
elementCount Int
rows Int
cols
    Right (Matrix rows cols (replicate count 0))

{- | Construct a square matrix with the finite input list on its diagonal and
zeros elsewhere. An empty list produces a @0 x 0@ matrix.
-}
diagonal :: (Num a) => [a] -> Maybe (Matrix a)
diagonal :: forall a. Num a => [a] -> Maybe (Matrix a)
diagonal = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Either MathError (Matrix a) -> Maybe (Matrix a))
-> ([a] -> Either MathError (Matrix a)) -> [a] -> Maybe (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Either MathError (Matrix a)
forall a. Num a => [a] -> Either MathError (Matrix a)
diagonalChecked

-- | Checked 'diagonal'. Rejects a square shape whose element count overflows.
diagonalChecked :: (Num a) => [a] -> Either MathError (Matrix a)
diagonalChecked :: forall a. Num a => [a] -> Either MathError (Matrix a)
diagonalChecked [a]
values = do
    let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
values
    _ <- Int -> Int -> Either MathError Int
elementCount Int
n Int
n
    Right (Matrix n n (concatMap (diagonalRow n) (zip [0 ..] values)))

-- | Look up an element using zero-based row and column indices.
at :: Matrix a -> Int -> Int -> Maybe a
at :: forall a. Matrix a -> Int -> Int -> Maybe a
at Matrix a
matrix Int
row Int
col = Either MathError a -> Maybe a
forall e a. Either e a -> Maybe a
toMaybe (Matrix a -> Int -> Int -> Either MathError a
forall a. Matrix a -> Int -> Int -> Either MathError a
atChecked Matrix a
matrix Int
row Int
col)

-- | Checked 'at'. Empty matrices have no valid element indices.
atChecked :: Matrix a -> Int -> Int -> Either MathError a
atChecked :: forall a. Matrix a -> Int -> Int -> Either MathError a
atChecked (Matrix Int
rows Int
cols [a]
values) Int
row Int
col
    | Int
row Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
col Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
row Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
rows Bool -> Bool -> Bool
|| Int
col Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
cols = Either MathError a
forall {b}. Either MathError b
invalid
    | Bool
otherwise = case Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
row Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
cols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
col) [a]
values of
        a
value : [a]
_ -> a -> Either MathError a
forall a b. b -> Either a b
Right a
value
        [] -> Either MathError a
forall {b}. Either MathError b
invalid
  where
    invalid :: Either MathError b
invalid = MathError -> Either MathError b
forall a b. a -> Either a b
Left ((Int, Int) -> (Int, Int) -> MathError
IndexOutOfBounds (Int
rows, Int
cols) (Int
row, Int
col))

{- | Extract a slice given its @(startRow, startColumn)@ and
@(rowCount, columnCount)@. Extents must be nonnegative and contained within
the matrix. A zero extent may start at the corresponding shape endpoint.
-}
sliceMatrix :: (Int, Int) -> (Int, Int) -> Matrix a -> Maybe (Matrix a)
sliceMatrix :: forall a. (Int, Int) -> (Int, Int) -> Matrix a -> Maybe (Matrix a)
sliceMatrix (Int, Int)
start (Int, Int)
extent = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Either MathError (Matrix a) -> Maybe (Matrix a))
-> (Matrix a -> Either MathError (Matrix a))
-> Matrix a
-> Maybe (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> (Int, Int) -> Matrix a -> Either MathError (Matrix a)
forall a.
(Int, Int) -> (Int, Int) -> Matrix a -> Either MathError (Matrix a)
sliceMatrixChecked (Int, Int)
start (Int, Int)
extent

{- | Checked 'sliceMatrix'. Extent shapes are validated before bounds; bounds
checks avoid adding the start and extent, which could overflow.
-}
sliceMatrixChecked ::
    (Int, Int) -> (Int, Int) -> Matrix a -> Either MathError (Matrix a)
sliceMatrixChecked :: forall a.
(Int, Int) -> (Int, Int) -> Matrix a -> Either MathError (Matrix a)
sliceMatrixChecked
    start :: (Int, Int)
start@(Int
startRow, Int
startCol)
    extent :: (Int, Int)
extent@(Int
height, Int
width)
    (Matrix Int
rows Int
cols [a]
values) = do
        count <- Int -> Int -> Either MathError Int
elementCount Int
height Int
width
        if startRow < 0
            || startCol < 0
            || startRow > rows
            || startCol > cols
            || height > rows - startRow
            || width > cols - startCol
            then Left (SliceOutOfBounds (rows, cols) start extent)
            else Right (Matrix height width (sliceValues count))
      where
        sliceValues :: a -> [a]
sliceValues a
0 = []
        sliceValues a
_ =
            ([a] -> [a]) -> [[a]] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
                (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
width ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
startCol)
                (Int -> [[a]] -> [[a]]
forall a. Int -> [a] -> [a]
take Int
height (Int -> [[a]] -> [[a]]
forall a. Int -> [a] -> [a]
drop Int
startRow (Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
cols [a]
values)))

{- | Change the shape while preserving row-major element order. The new shape
must contain exactly the same number of elements.
-}
reshape :: Int -> Int -> Matrix a -> Maybe (Matrix a)
reshape :: forall a. Int -> Int -> Matrix a -> Maybe (Matrix a)
reshape Int
rows Int
cols = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Either MathError (Matrix a) -> Maybe (Matrix a))
-> (Matrix a -> Either MathError (Matrix a))
-> Matrix a
-> Maybe (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Matrix a -> Either MathError (Matrix a)
forall a. Int -> Int -> Matrix a -> Either MathError (Matrix a)
reshapeChecked Int
rows Int
cols

-- | Checked 'reshape'. Validate the target shape before comparing counts.
reshapeChecked :: Int -> Int -> Matrix a -> Either MathError (Matrix a)
reshapeChecked :: forall a. Int -> Int -> Matrix a -> Either MathError (Matrix a)
reshapeChecked Int
rows Int
cols (Matrix Int
oldRows Int
oldCols [a]
values) = do
    expected <- Int -> Int -> Either MathError Int
elementCount Int
rows Int
cols
    let actual = Int
oldRows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
oldCols
    if expected == actual
        then Right (Matrix rows cols values)
        else Left (ElementCountMismatch expected actual)

-- | Add matrices with equal shapes.
addMatrix :: (Num a) => Matrix a -> Matrix a -> Maybe (Matrix a)
addMatrix :: forall a. Num a => Matrix a -> Matrix a -> Maybe (Matrix a)
addMatrix Matrix a
a Matrix a
b = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
addMatrixChecked Matrix a
a Matrix a
b)

-- | Checked 'addMatrix'.
addMatrixChecked :: (Num a) => Matrix a -> Matrix a -> Either MathError (Matrix a)
addMatrixChecked :: forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
addMatrixChecked = (a -> a -> a)
-> Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a b c.
(a -> b -> c)
-> Matrix a -> Matrix b -> Either MathError (Matrix c)
zipMatrixChecked a -> a -> a
forall a. Num a => a -> a -> a
(+)

-- | Subtract matrices with equal shapes.
subtractMatrix :: (Num a) => Matrix a -> Matrix a -> Maybe (Matrix a)
subtractMatrix :: forall a. Num a => Matrix a -> Matrix a -> Maybe (Matrix a)
subtractMatrix Matrix a
a Matrix a
b = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
subtractMatrixChecked Matrix a
a Matrix a
b)

-- | Checked 'subtractMatrix'.
subtractMatrixChecked ::
    (Num a) => Matrix a -> Matrix a -> Either MathError (Matrix a)
subtractMatrixChecked :: forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
subtractMatrixChecked = (a -> a -> a)
-> Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a b c.
(a -> b -> c)
-> Matrix a -> Matrix b -> Either MathError (Matrix c)
zipMatrixChecked (-)

-- | Multiply corresponding elements of matrices with equal shapes.
hadamard :: (Num a) => Matrix a -> Matrix a -> Maybe (Matrix a)
hadamard :: forall a. Num a => Matrix a -> Matrix a -> Maybe (Matrix a)
hadamard Matrix a
a Matrix a
b = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
hadamardChecked Matrix a
a Matrix a
b)

-- | Checked 'hadamard'.
hadamardChecked :: (Num a) => Matrix a -> Matrix a -> Either MathError (Matrix a)
hadamardChecked :: forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
hadamardChecked = (a -> a -> a)
-> Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a b c.
(a -> b -> c)
-> Matrix a -> Matrix b -> Either MathError (Matrix c)
zipMatrixChecked a -> a -> a
forall a. Num a => a -> a -> a
(*)

-- | Multiply every element by the scalar, preserving the shape.
scaleMatrix :: (Num a) => a -> Matrix a -> Matrix a
scaleMatrix :: forall a. Num a => a -> Matrix a -> Matrix a
scaleMatrix a
scalar (Matrix Int
rows Int
cols [a]
values) =
    Int -> Int -> [a] -> Matrix a
forall a. Int -> Int -> [a] -> Matrix a
Matrix Int
rows Int
cols ((a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
scalar a -> a -> a
forall a. Num a => a -> a -> a
*) [a]
values)

{- | Swap rows and columns. This preserves zero dimensions, so transposing
twice returns the original matrix even when one dimension is zero.
-}
transposeMatrix :: Matrix a -> Matrix a
transposeMatrix :: forall a. Matrix a -> Matrix a
transposeMatrix (Matrix Int
rows Int
cols [a]
values)
    | Int
rows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
cols Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> Int -> [a] -> Matrix a
forall a. Int -> Int -> [a] -> Matrix a
Matrix Int
cols Int
rows []
    | Bool
otherwise = Int -> Int -> [a] -> Matrix a
forall a. Int -> Int -> [a] -> Matrix a
Matrix Int
cols Int
rows ([[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose (Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
cols [a]
values)))

{- | Matrix product when the left column count equals the right row count and
the output element count fits in 'Int'. A product with zero inner dimension
contains zeros, with the full output shape preserved.
-}
multiply :: (Num a) => Matrix a -> Matrix a -> Maybe (Matrix a)
multiply :: forall a. Num a => Matrix a -> Matrix a -> Maybe (Matrix a)
multiply Matrix a
a Matrix a
b = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Matrix a -> Matrix a -> Either MathError (Matrix a)
forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
multiplyChecked Matrix a
a Matrix a
b)

-- | Checked 'multiply'. Inner dimensions are checked before the output shape.
multiplyChecked :: (Num a) => Matrix a -> Matrix a -> Either MathError (Matrix a)
multiplyChecked :: forall a.
Num a =>
Matrix a -> Matrix a -> Either MathError (Matrix a)
multiplyChecked (Matrix Int
r1 Int
c1 [a]
a) (Matrix Int
r2 Int
c2 [a]
b)
    | Int
c1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r2 = MathError -> Either MathError (Matrix a)
forall a b. a -> Either a b
Left (Int -> Int -> MathError
InnerDimensionMismatch Int
c1 Int
r2)
    | Bool
otherwise = do
        count <- Int -> Int -> Either MathError Int
elementCount Int
r1 Int
c2
        Right (Matrix r1 c2 (result count))
  where
    -- Empty outputs and zero inner dimensions must not enumerate empty
    -- input rows; their dimension may be as large as maxBound :: Int.
    result :: Int -> [a]
result Int
0 = []
    result Int
count
        | Int
c1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
count a
0
        | Bool
otherwise =
            [ [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
row [a]
col)
            | [a]
row <- Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
c1 [a]
a
            , [a]
col <- [[a]]
columns
            ]
    columns :: [[a]]
columns = [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
transpose (Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
c2 [a]
b)

{- | Construct a square identity matrix of nonnegative order. The element
count must fit in 'Int'.
-}
identity :: (Num a) => Int -> Maybe (Matrix a)
identity :: forall a. Num a => Int -> Maybe (Matrix a)
identity = Either MathError (Matrix a) -> Maybe (Matrix a)
forall e a. Either e a -> Maybe a
toMaybe (Either MathError (Matrix a) -> Maybe (Matrix a))
-> (Int -> Either MathError (Matrix a)) -> Int -> Maybe (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either MathError (Matrix a)
forall a. Num a => Int -> Either MathError (Matrix a)
identityChecked

-- | Checked 'identity'.
identityChecked :: (Num a) => Int -> Either MathError (Matrix a)
identityChecked :: forall a. Num a => Int -> Either MathError (Matrix a)
identityChecked Int
n = do
    _ <- Int -> Int -> Either MathError Int
elementCount Int
n Int
n
    Right (Matrix n n (concatMap (\Int
i -> Int -> (Int, a) -> [a]
forall a. Num a => Int -> (Int, a) -> [a]
diagonalRow Int
n (Int
i, a
1)) [0 .. n - 1]))

-- Validate with unbounded arithmetic before any Int multiplication or storage
-- allocation. Every constructor and shape-changing operation uses this gate.
elementCount :: Int -> Int -> Either MathError Int
elementCount :: Int -> Int -> Either MathError Int
elementCount Int
rows Int
cols
    | Int
rows Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
cols Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Int
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
NegativeDimensions (Int
rows, Int
cols))
    | Integer
count Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) =
        MathError -> Either MathError Int
forall a b. a -> Either a b
Left ((Int, Int) -> MathError
ElementCountOverflow (Int
rows, Int
cols))
    | Bool
otherwise = Int -> Either MathError Int
forall a b. b -> Either a b
Right (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
count)
  where
    count :: Integer
count = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
rows Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
cols

zipMatrixChecked ::
    (a -> b -> c) -> Matrix a -> Matrix b -> Either MathError (Matrix c)
zipMatrixChecked :: forall a b c.
(a -> b -> c)
-> Matrix a -> Matrix b -> Either MathError (Matrix c)
zipMatrixChecked a -> b -> c
operation (Matrix Int
r1 Int
c1 [a]
a) (Matrix Int
r2 Int
c2 [b]
b)
    | Int
r1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r2 Bool -> Bool -> Bool
|| Int
c1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c2 = MathError -> Either MathError (Matrix c)
forall a b. a -> Either a b
Left ((Int, Int) -> (Int, Int) -> MathError
ShapeMismatch (Int
r1, Int
c1) (Int
r2, Int
c2))
    | Bool
otherwise = Matrix c -> Either MathError (Matrix c)
forall a b. b -> Either a b
Right (Int -> Int -> [c] -> Matrix c
forall a. Int -> Int -> [a] -> Matrix a
Matrix Int
r1 Int
c1 ((a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> b -> c
operation [a]
a [b]
b))

diagonalRow :: (Num a) => Int -> (Int, a) -> [a]
diagonalRow :: forall a. Num a => Int -> (Int, a) -> [a]
diagonalRow Int
n (Int
index, a
value) =
    Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
index a
0 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ a
value a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a
0

-- Only positive widths are needed for stored elements. Handling zero here
-- explicitly also prevents a non-consuming recursion for empty shapes.
chunks :: Int -> [a] -> [[a]]
chunks :: forall a. Int -> [a] -> [[a]]
chunks Int
width [a]
values
    | Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = []
    | Bool
otherwise = case [a]
values of
        [] -> []
        [a]
_ ->
            let ([a]
row, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
width [a]
values
             in [a]
row [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
width [a]
rest

toMaybe :: Either e a -> Maybe a
toMaybe :: forall e a. Either e a -> Maybe a
toMaybe = (e -> Maybe a) -> (a -> Maybe a) -> Either e a -> Maybe a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe a -> e -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) a -> Maybe a
forall a. a -> Maybe a
Just