mldsa-0.1.0.0: src/Matrix.hs

-- |
-- Module      : Matrix
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- A matrix here is simply a vector of vectors.  The module also implements two
-- utility functions 'mmul' and 'mmulAdd' that multiply a matrix and a vector.
--
module Matrix
    ( create, mmul, mmulAdd
    ) where

import Base
import Math
import Vector (Vector)
import qualified Vector

create :: (KnownNat m, KnownNat n) => (Offset ty -> Offset (Vector n ty) -> ty) -> Vector m (Vector n ty)
create f = Vector.create $ \j -> Vector.create (`f` j)
{-# INLINE create #-}

index :: Vector m (Vector n ty) -> Offset (Vector n ty) -> Offset ty -> ty
index a i = Vector.index (Vector.index a i)

mmul :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a
mmul a u = Vector.seq u (fmap (`Vector.dot` u) a)
{-# INLINE mmul #-}

mmulAdd :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a -> Vector m a
mmulAdd a u b = Vector.seq u (Vector.mapIx f b)
    where
        f (Offset i) bv =
            let g (Offset j) vu = (index a (Offset i) (Offset j), vu)
             in Vector.biMulFoldIndexWith g bv u
{-# INLINE mmulAdd #-}