-- |
-- Module      : Matrix
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- A matrix here is simply a vector of vectors.  The module also implements
-- 'mulz' as dot product and two utility functions 'mulw' and 'muly' that
-- multiply a matrix and a vector.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
module Matrix
    ( create, mulw, muly, mulz
#ifdef ML_KEM_TESTING
    , transpose
#endif
    ) 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 :: forall (m :: Nat) (n :: Nat) ty.
(KnownNat m, KnownNat n) =>
(Offset ty -> Offset (Vector n ty) -> ty) -> Vector m (Vector n ty)
create Offset ty -> Offset (Vector n ty) -> ty
f = (Offset (Vector n ty) -> Vector n ty) -> Vector m (Vector n ty)
forall (n :: Nat) a. KnownNat n => (Offset a -> a) -> Vector n a
Vector.create ((Offset (Vector n ty) -> Vector n ty) -> Vector m (Vector n ty))
-> (Offset (Vector n ty) -> Vector n ty) -> Vector m (Vector n ty)
forall a b. (a -> b) -> a -> b
$ \Offset (Vector n ty)
j -> (Offset ty -> ty) -> Vector n ty
forall (n :: Nat) a. KnownNat n => (Offset a -> a) -> Vector n a
Vector.create (Offset ty -> Offset (Vector n ty) -> ty
`f` Offset (Vector n ty)
j)
{-# INLINE create #-}

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

mulw :: (KnownNat n, BiMulAdd b a) => Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a
mulw :: forall (n :: Nat) b a (m :: Nat).
(KnownNat n, BiMulAdd b a) =>
Vector m (Vector n b) -> Vector m a -> Vector n a -> Vector n a
mulw Vector m (Vector n b)
a !Vector m a
u !Vector n a
b = (Offset a -> a) -> Vector n a
forall (n :: Nat) a. KnownNat n => (Offset a -> a) -> Vector n a
Vector.create ((Offset a -> a) -> Vector n a) -> (Offset a -> a) -> Vector n a
forall a b. (a -> b) -> a -> b
$ \(Offset Int
i) ->
    (a -> Offset a -> a -> a) -> a -> Vector m a -> a
forall c a (n :: Nat).
(c -> Offset a -> a -> c) -> c -> Vector n a -> c
Vector.foldIndexWith (\a
c (Offset Int
j) a
vu -> b -> a -> a -> a
forall b a. BiMulAdd b a => b -> a -> a -> a
biMulAdd (Vector m (Vector n b) -> Offset b -> Offset (Vector n b) -> b
forall (m :: Nat) (n :: Nat) ty.
Vector m (Vector n ty) -> Offset ty -> Offset (Vector n ty) -> ty
index Vector m (Vector n b)
a (Int -> Offset b
forall ty. Int -> Offset ty
Offset Int
i) (Int -> Offset (Vector n b)
forall ty. Int -> Offset ty
Offset Int
j)) a
vu a
c) (Vector n a -> Offset a -> a
forall (n :: Nat) a. Vector n a -> Offset a -> a
Vector.index Vector n a
b (Int -> Offset a
forall ty. Int -> Offset ty
Offset Int
i)) Vector m a
u

muly :: BiMulAdd b a => Vector m (Vector n b) -> Vector n a -> Vector m a
muly :: forall b a (m :: Nat) (n :: Nat).
BiMulAdd b a =>
Vector m (Vector n b) -> Vector n a -> Vector m a
muly Vector m (Vector n b)
a !Vector n a
u = (Vector n b -> a) -> Vector m (Vector n b) -> Vector m a
forall a b. (a -> b) -> Vector m a -> Vector m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vector n b -> Vector n a -> a
forall b a (n :: Nat).
BiMulAdd b a =>
Vector n b -> Vector n a -> a
`mulz` Vector n a
u) Vector m (Vector n b)
a

mulz :: BiMulAdd b a => Vector n b -> Vector n a -> a
mulz :: forall b a (n :: Nat).
BiMulAdd b a =>
Vector n b -> Vector n a -> a
mulz = (a -> b -> a -> a)
-> (b -> a -> a) -> Vector n b -> Vector n a -> a
forall c a b (n :: Nat).
(c -> a -> b -> c)
-> (a -> b -> c) -> Vector n a -> Vector n b -> c
Vector.fold1ZipWith (\a
c b
a a
b -> b -> a -> a -> a
forall b a. BiMulAdd b a => b -> a -> a -> a
biMulAdd b
a a
b a
c) b -> a -> a
forall b a. BiMul b a => b -> a -> a
(..*)

#ifdef ML_KEM_TESTING
transpose :: (KnownNat m, KnownNat n) => Vector m (Vector n ty) -> Vector n (Vector m ty)
transpose a = create $ \(Offset j) (Offset i) -> index a (Offset i) (Offset j)
#endif