mldsa-0.1.0.0: src/Vector.hs

-- |
-- Module      : Vector
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- A vector of lifted elements with the vector dimension at type level.
-- Backed by type t'SmallArray' from primitive.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Vector
    ( Vector, Vector.concatMap, Vector.dot, Vector.index, Vector.toNormalForm
    , Vector.create, Vector.createMaybe, Vector.createMaybeAccum, Vector.foldl'
    , Vector.fold1ZipWith, Vector.biMulFoldIndexWith, mapIx, Vector.seq
    , Vector.unzipWith, Vector.zipWith
#ifdef ML_DSA_TESTING
    , Vector.replicateM
#endif
    ) where

import Data.Primitive.SmallArray

import Control.DeepSeq (NFData(..))
#ifdef ML_DSA_TESTING
import Control.Monad
#endif
import Control.Monad.ST

#if !(MIN_VERSION_base(4,20,0))
import Data.List as Prelude (foldl')
#endif
import Data.Proxy

import Base
import Iterate
import Math

type Array = SmallArray
type MArray ty s = SmallMutableArray s ty

newtype Vector (n :: Nat) a = Vector { unVector :: Array a }
    deriving (Eq, Show)

instance Functor (Vector n) where
    fmap = mapVector
    {-# INLINE fmap #-}

instance (Add a, KnownNat n) => Add (Vector n a) where
    zero = create (const zero)
    {-# INLINE zero #-}
    (.+) = Vector.zipWith (.+)
    {-# INLINE (.+) #-}
    (.-) = Vector.zipWith (.-)
    {-# INLINE (.-) #-}
    neg = mapVector neg
    {-# INLINE neg #-}

arrayCreate :: forall ty. CountOf ty -> (Offset ty -> ty) -> Array ty
arrayCreate n initializer = runST (arrayNew n >>= iter initializer)
  where
    iter :: PrimMonad prim => (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)
    iter f ma = loop 0
      where
        loop s@(Offset i)
            | s .==# n = unsafeFreezeSmallArray ma
            | otherwise = writeSmallArray ma i (f s) >> loop (s + 1)
        {-# INLINE loop #-}
    {-# INLINE iter #-}

arrayLength :: Array ty -> CountOf ty
arrayLength = CountOf . sizeofSmallArray

arrayMapIx :: (Offset a -> a -> b) -> Array a -> Array b
arrayMapIx f a = arrayCreate (CountOf sz) $ \(Offset i) ->
    let off = Offset i in f off (arrayIndex a off)
  where CountOf sz = arrayLength a

arrayNew :: PrimMonad prim => CountOf ty -> prim (MArray ty (PrimState prim))
arrayNew (CountOf c) = newSmallArray c placeholder
  where placeholder = error "arrayNew: unexpected evaluation"

create :: forall n a. KnownNat n => (Offset a -> a) -> Vector n a
create f = Vector $ arrayCreate (CountOf sz) (\(Offset !i) -> f (Offset i))
  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)
{-# INLINE [2] create #-}

mapIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b
mapIx = mapVectorIx
{-# INLINE [2] mapIx #-}

mapVector :: (a -> b) -> Vector n a -> Vector n b
mapVector f = mapVectorIx $ \_ x -> f x
{-# INLINE [2] mapVector #-}

mapVectorIx :: (Offset a -> a -> b) -> Vector n a -> Vector n b
mapVectorIx f = Vector <$> arrayMapIx f . unVector
{-# INLINE [1] mapVectorIx #-}

arrayIndex :: Array a -> Offset a -> a
#ifdef ML_DSA_TESTING
arrayIndex a off@(Offset i) =
    checkBounds (arrayLength a) off $ indexSmallArray a i

replicateM :: forall n m a. (KnownNat n, Applicative m) => m a -> m (Vector n a)
replicateM f = Vector . smallArrayFromList <$> Control.Monad.replicateM sz f
  where !sz = fromIntegral $ natVal (Proxy :: Proxy n)
#else
arrayIndex a (Offset i) = indexSmallArray a i
#endif

index :: Vector n a -> Offset a -> a
index = arrayIndex . unVector

concatMap :: Monoid b => (a -> b) -> Vector n a -> b
concatMap f = mconcat . mapToList f
{-# INLINE concatMap #-}

mapToList :: (a -> b) -> Vector n a -> [b]
mapToList f (Vector a) = Prelude.map (f . arrayIndex a . Offset) (offsets sa)
  where CountOf sa = arrayLength a

foldl' :: (b -> a -> b) -> b -> Vector n a -> b
foldl' f b (Vector a) = Prelude.foldl' g b (offsets sa)
  where
    g acc i = acc `f` arrayIndex a (Offset i)
    CountOf !sa = arrayLength a
{-# INLINE foldl' #-}

seq :: b -> Vector n a -> Vector n a
seq = Prelude.seq
{-# INLINE [1] seq #-}

zipWith :: (a -> b -> c) -> Vector n a -> Vector n b -> Vector n c
zipWith f a (Vector !b) = mapVectorIx g a
  where g (Offset i) x = f x $ arrayIndex b (Offset i)
{-# INLINE [2] zipWith #-}

fold1ZipWith :: (c -> a -> b -> c) -> (a -> b -> c) -> Vector n a -> Vector n b -> c
fold1ZipWith f g (Vector a) (Vector !b) =
    Prelude.foldl' ff gg (offsetsFrom 1 sa)
  where
    ff x i = f x (arrayIndex a (Offset i)) (arrayIndex b (Offset i))
    gg = g (arrayIndex a 0) (arrayIndex b 0)
    CountOf !sa = arrayLength a
{-# INLINE fold1ZipWith #-}

biMulFoldIndexWith :: BiMulAdd b a => (Offset ty -> t -> (b, a)) -> a -> Vector n t -> a
biMulFoldIndexWith f c (Vector a) =
    biMulFold c (map g $ offsets sa)
  where
    g i = f (Offset i) (arrayIndex a (Offset i))
    CountOf !sa = arrayLength a
{-# INLINE biMulFoldIndexWith #-}

dot :: BiMulAdd b a => Vector n b -> Vector n a -> a
dot (Vector b) (Vector a) =
    biMulFold (arrayIndex b 0 ..* arrayIndex a 0) (map g $ offsetsFrom 1 sb)
  where
    g i = (arrayIndex b (Offset i), arrayIndex a (Offset i))
    CountOf !sb = arrayLength b
{-# INLINE dot #-}

toNormalForm :: NFData a => Vector n a -> ()
toNormalForm = Vector.foldl' (\acc x -> acc `Prelude.seq` rnf x) ()

createMaybe :: KnownNat n => (Offset a -> Maybe a) -> Maybe (Vector n a)
createMaybe f = createMaybeAccum (\_ off -> ((), ) <$> f off) (const Just) ()
{-# INLINE createMaybe #-}

createMaybeAccum :: forall n a b r. KnownNat n => (b -> Offset a -> Maybe (b, a)) -> (b -> Vector n a -> Maybe r) -> b -> Maybe r
createMaybeAccum f k b0 = runST $ arrayNew (CountOf n) >>= \ma -> loop ma b0 0
  where
    !n = fromIntegral $ natVal (Proxy :: Proxy n)

    loop :: MArray a s -> b -> Offset a -> ST s (Maybe r)
    loop !ma b s@(Offset i)
        | s .==# CountOf n =
            unsafeFreezeSmallArray ma >>= \a -> return (b `k` Vector a)
        | otherwise = do
            case f b s of
                Nothing -> return Nothing
                Just (b', a) -> do
                    writeSmallArray ma i a
                    loop ma b' (s + 1)
{-# INLINE createMaybeAccum #-}

unzipWith :: forall n a b c. KnownNat n => (a -> (b, c)) -> Vector n a -> (Vector n b, Vector n c)
unzipWith f (Vector !a) = runST $ do
    mb <- arrayNew (CountOf n)
    mc <- arrayNew (CountOf n)
    loop mb mc 0
  where
    !n = fromIntegral $ natVal (Proxy :: Proxy n)

    loop :: MArray b s -> MArray c s -> Offset a -> ST s (Vector n b, Vector n c)
    loop !mb !mc s@(Offset i)
        | s .==# CountOf n = do
            b <- unsafeFreezeSmallArray mb
            c <- unsafeFreezeSmallArray mc
            return (Vector b, Vector c)
        | otherwise = do
            let (b, c) = f (arrayIndex a s)
            writeSmallArray mb i b
            writeSmallArray mc i c
            loop mb mc (s + 1)
{-# INLINE unzipWith #-}


-- Rewrite rules
--
-- A first set of rules before Phase 2 performs simplifications between the four
-- main functions: create, mapVector, mapIx, and zipWith.
--
-- During Phase 2, the functions are then inlined.  While function create
-- is replaced with its final form as call to arrayCreate, the three other
-- functions mapVector, mapIx, and zipWith all become calls to mapVectorIx.
-- Then, nested calls to mapVectorIx can further be simplified using rule
-- "mapVectorIx/mapVectorIx".
--
-- Finally at Phase 1 the function mapVectorIx is replaced with a call to
-- arrayCreate.
--
-- Both layers of rules have transformations that push calls to Vector.seq
-- outwards.  Normal 'Prelude.seq' or the use of bang patterns would prevent
-- rewrite rules from firing.

{-# RULES
"mapVector/mapVector" [~2] forall f g a. mapVector f (mapVector g a) = mapVector (f . g) a
"mapVector/mapIx" [~2] forall f g a. mapVector f (mapIx g a) = mapIx (\i -> f . g i) a
"mapVector/create" [~2] forall f g. mapVector f (create g) = create (\(Offset i) -> f (g (Offset i)))
"mapVector/zipWith" [~2] forall f g a b. mapVector f (Vector.zipWith g a b) = Vector.zipWith (\x -> f . g x) a b
"mapVector/seq" [~2] forall f a b. mapVector f (Vector.seq b a) = Vector.seq b (mapVector f a)

"zipWith/mapVector left" [~2] forall f g a. Vector.zipWith f (mapVector g a) = Vector.zipWith (f . g) a
"zipWith/mapVector right" [~2] forall f g a b. Vector.zipWith f a (mapVector g b) = Vector.zipWith (\aa bb -> f aa (g bb)) a b
"zipWith/create left" [~2] forall f g. Vector.zipWith f (create g) = mapIx (\(Offset i) -> f (g (Offset i)))
"zipWith/create right" [~2] forall f g a. Vector.zipWith f a (create g) = mapIx (\(Offset i) x -> f x (g (Offset i))) a
"zipWith/zipWith right" [~2] forall f g a b c. Vector.zipWith f a (Vector.zipWith g b c)  = Vector.zipWith (flip f) (Vector.zipWith g b c) a
"zipWith/seq left" [~2] forall f a b c. Vector.zipWith f (Vector.seq c a) b = Vector.seq c (Vector.zipWith f a b)
"zipWith/seq right" [~2] forall f a b c. Vector.zipWith f a (Vector.seq c b) = Vector.seq c (Vector.zipWith f a b)

"mapIx/mapVector" [~2] forall f g a. mapIx f (mapVector g a) = mapIx (\(Offset i) -> f (Offset i) . g) a
"mapIx/mapIx" [~2] forall f g a. mapIx f (mapIx g a) = mapIx (\(Offset i) -> f (Offset i) . g (Offset i)) a
"mapIx/create" [~2] forall f g. mapIx f (create g) = create (\(Offset i) -> f (Offset i) (g (Offset i)))
"mapIx/seq" [~2] forall f a b. mapIx f (Vector.seq b a) = Vector.seq b (mapIx f a)

"mapVectorIx/mapVectorIx" [~1] forall f g a. mapVectorIx f (mapVectorIx g a) = mapVectorIx (\(Offset i) -> f (Offset i) . g (Offset i)) a
"mapVectorIx/seq" [~1] forall f a b. mapVectorIx f (Vector.seq b a) = Vector.seq b (mapVectorIx f a)
  #-}