mldsa-0.1.0.0: src/Auxiliary.hs

-- |
-- Module      : Auxiliary
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- ML-DSA auxiliary functions
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
module Auxiliary
    ( Zq, Rq, Tq
    , Norm(..), normFrom, countFrom, ntt, nttInv
    , simpleBitPack, simpleBitPack10, simpleBitUnpack10
    , bitPack, bitUnpack, bitPackSafe, bitUnpackSafe
    , hintBitPack, hintBitUnpack
    , sampleInBall, rejNttPoly, rejBoundedPoly
    , powerTwoRound
    , highBits, lowBits
    , Hints, makeHint, useHint
#ifdef ML_DSA_TESTING
    , Mq, toMontgomery, fromMontgomery, norm
    , bitRev8, fromZq, toZq, fromCoeffs, toCoeffs, fromBools, toBools
    , powerTwoRoundZq, decomposeZq
    , simpleBitUnpack
#endif
    ) where

import Crypto.Hash.Algorithms

import Data.ByteArray (ByteArrayAccess)
import qualified Data.ByteArray as B
import qualified Data.Memory.Endian as B

import Data.Primitive.Types (Prim(..))

import Control.DeepSeq (NFData(..))
import Control.Monad
import Control.Monad.ST

import Data.Bits
import Data.Int
import Data.Proxy
import Data.Word

#if MIN_VERSION_base(4,16,0)
import GHC.Exts as Exts (eqWord32#, int2Word#, wordToWord8#)
#else
import GHC.Exts as Exts (eqWord#, int2Word#)
#endif
import GHC.TypeNats
import GHC.Word (Word8(W8#), Word32(W32#))

import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (pokeByteOff)

import Base
import Block (blockIndex)
import BlockN (BlockN, MutableBlockN)
import Builder (Builder)
import Crypto (BlockDigest, Digest)
import Fusion
import Machine
import Marking (Classified, SecurityMarking(..))
import SecureBlock (SecureBlock)
import SecureBytes (SecureBytes)
import Vector (Vector)
import qualified BlockN
import qualified Builder
import qualified ByteArrayST as ST
import qualified Crypto
import qualified Vector
import Math

type N = 256

n :: Int
n = 256

q :: Integer
q = 8380417

q32 :: Word32
q32 = fromInteger q

q64 :: Word64
q64 = fromInteger q

bitRev8 :: Word8 -> Word8
bitRev8 = fromIntegral . rev8 . fromIntegral
  where
    rev1, rev2, rev4, rev8 :: Word -> Word
    rev1 b = b .&. 1
    rev2 b = rev1 (b `unsafeShiftR` 1) .|. rev1 b `unsafeShiftL` 1
    rev4 b = rev2 (b `unsafeShiftR` 2) .|. rev2 b `unsafeShiftL` 2
    rev8 b = rev4 (b `unsafeShiftR` 4) .|. rev4 b `unsafeShiftL` 4

unsafeShiftIR :: Word32 -> Int -> Word32
unsafeShiftIR x s = fromIntegral ((fromIntegral x :: Int32) `unsafeShiftR` s)
{-# INLINE unsafeShiftIR #-}

select32 :: Word32 -> Word32 -> Word32 -> Word32
select32 mask yes no = (mask .&. yes) .|. (complement mask .&. no)
{-# INLINE select32 #-}

max32 :: Word32 -> Word32 -> Word32
max32 a b = let cond = (b - a) `unsafeShiftIR` 31 in select32 cond a b
{-# INLINE max32 #-}

min32 :: Word32 -> Word32 -> Word32
min32 a b = let cond = (a - b) `unsafeShiftIR` 31 in select32 cond a b
{-# INLINE min32 #-}

-- Reduction 𝑥 mod 5 for 0 ≤ 𝑥 < 15
mod5 :: Word32 -> Word32
mod5 x = let quotient = (x * 205) `unsafeShiftR` 10 in x - 5 * quotient

-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 2𝑞
reduceSimple :: Word32 -> Word32
reduceSimple x = select32 mask x subtracted
  where
    subtracted = x - q32
    mask = subtracted `unsafeShiftIR` 31
{-# INLINE reduceSimple #-}

#ifdef ML_DSA_TESTING
-- Reduction 𝑥 mod 𝑞 for 0 ≤ 𝑥 < 1025𝑞
reduce :: Word64 -> Word32
reduce x = reduceSimple $ fromIntegral (x - t * q64)
  where t = x `unsafeShiftR` 23
{-# INLINE reduce #-}
#endif

-- Computes 𝑎 . 2^-32 mod 𝑞
reduceMontgomery :: Word64 -> Word32
reduceMontgomery a = reduceSimple (fromIntegral r)
  where
    t = (fromIntegral a :: Word32) * 4236238847  -- inverse of -𝑞 modulo 2^32
    r = (a + fromIntegral t * q64) `unsafeShiftR` 32
{-# INLINE reduceMontgomery #-}

newtype Zq = Zq Word32
#ifdef ML_DSA_TESTING
    deriving (Eq, Show)
#else
    deriving Eq
#endif

instance Prim Zq where
    sizeOf# (Zq a) = sizeOf# a
    {-# INLINE sizeOf# #-}
    alignment# (Zq a) = alignment# a
    {-# INLINE alignment# #-}
#if MIN_VERSION_primitive(0,9,0)
    sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32)
    {-# INLINE sizeOfType# #-}
    alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32)
    {-# INLINE alignmentOfType# #-}
#endif
    indexByteArray# ba i = Zq (indexByteArray# ba i)
    {-# INLINE indexByteArray# #-}
    readByteArray# mba i s =
        case readByteArray# mba i s of
            (# s', a #) -> (# s', Zq a #)
    {-# INLINE readByteArray# #-}
    writeByteArray# mba i (Zq a) = writeByteArray# mba i a
    {-# INLINE writeByteArray# #-}
    setByteArray# mba i len (Zq a) = setByteArray# mba i len a
    {-# INLINE setByteArray# #-}
    indexOffAddr# addr i = Zq (indexOffAddr# addr i)
    {-# INLINE indexOffAddr# #-}
    readOffAddr# addr i s =
        case readOffAddr# addr i s of
            (# s', a #) -> (# s', Zq a #)
    {-# INLINE readOffAddr# #-}
    writeOffAddr# addr i (Zq a) = writeOffAddr# addr i a
    {-# INLINE writeOffAddr# #-}
    setOffAddr# addr i len (Zq a) = setOffAddr# addr i len a
    {-# INLINE setOffAddr# #-}

instance PrimSized Zq where
    type PrimSize Zq = 4

instance Add Zq where
    zero = Zq 0
    Zq a .+ Zq b = Zq $ reduceSimple (a + b)
    Zq a .- Zq b = Zq $ reduceSimple (a + q32 - b)
    neg (Zq a) = Zq $ reduceSimple (q32 - a)

#ifdef ML_DSA_TESTING
instance Mul Zq where
    one = Zq 1
    a .* b = fromMontgomery (toMontgomery a .* toMontgomery b)

toZq :: Word32 -> Zq
toZq = Zq . reduce . fromIntegral
#endif

fromZq :: Zq -> Word32
fromZq (Zq a) = a

newtype Mq = Mq Word32
#ifdef ML_DSA_TESTING
    deriving (Eq, Show)
#else
    deriving Eq
#endif

instance Prim Mq where
    sizeOf# (Mq a) = sizeOf# a
    {-# INLINE sizeOf# #-}
    alignment# (Mq a) = alignment# a
    {-# INLINE alignment# #-}
#if MIN_VERSION_primitive(0,9,0)
    sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word32)
    {-# INLINE sizeOfType# #-}
    alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word32)
    {-# INLINE alignmentOfType# #-}
#endif
    indexByteArray# ba i = Mq (indexByteArray# ba i)
    {-# INLINE indexByteArray# #-}
    readByteArray# mba i s =
        case readByteArray# mba i s of
            (# s', a #) -> (# s', Mq a #)
    {-# INLINE readByteArray# #-}
    writeByteArray# mba i (Mq a) = writeByteArray# mba i a
    {-# INLINE writeByteArray# #-}
    setByteArray# mba i len (Mq a) = setByteArray# mba i len a
    {-# INLINE setByteArray# #-}
    indexOffAddr# addr i = Mq (indexOffAddr# addr i)
    {-# INLINE indexOffAddr# #-}
    readOffAddr# addr i s =
        case readOffAddr# addr i s of
            (# s', a #) -> (# s', Mq a #)
    {-# INLINE readOffAddr# #-}
    writeOffAddr# addr i (Mq a) = writeOffAddr# addr i a
    {-# INLINE writeOffAddr# #-}
    setOffAddr# addr i len (Mq a) = setOffAddr# addr i len a
    {-# INLINE setOffAddr# #-}

instance PrimSized Mq where
    type PrimSize Mq = 4

instance Add Mq where
    zero = Mq 0
    Mq a .+ Mq b = Mq $ reduceSimple (a + b)
    Mq a .- Mq b = Mq $ reduceSimple (a + q32 - b)
    neg (Mq a) = Mq $ reduceSimple (q32 - a)

instance Mul Mq where
    one = Mq 4193792  -- toMongomery (Zq 1)
    Mq a .* Mq b = Mq $ reduceMontgomery (fromIntegral a * fromIntegral b)

instance MulAdd Mq where
    mulAdd a b c = a .* b .+ c

#ifdef ML_DSA_TESTING
instance BiMul Mq Mq where
    (..*) = (.*)

instance BiMulAdd Mq Mq where
    biMulAdd = mulAdd
#endif

toMontgomery :: Zq -> Mq
toMontgomery (Zq a) = Mq $ reduceMontgomery $ fromIntegral a * 2365951  -- R^2 mod 𝑞

fromMontgomery :: Mq -> Zq
fromMontgomery (Mq m) = Zq $ reduceMontgomery $ fromIntegral m

newtype Rq marking = Rq (BlockN marking N Zq)
#ifdef ML_DSA_TESTING
    deriving (Eq, Show, NFData)
#else
    deriving NFData
#endif

instance Classified marking => Add (Rq marking) where
    zero = Rq zero
    Rq a .+ Rq b = Rq (a .+ b)
    {-# INLINE (.+) #-}
    Rq a .- Rq b = Rq (a .- b)
    {-# INLINE (.-) #-}
    neg (Rq a) = Rq (neg a)
    {-# INLINE neg #-}

instance Crypto.ConstEqW (Rq Sec) where
    constEqW (Rq a) (Rq b) = Crypto.constEqW
        (BlockN.unsafeCast a :: SecureBlock Sec Word)
        (BlockN.unsafeCast b :: SecureBlock Sec Word)

instance Crypto.ConstEqW (Rq Pub) where
    constEqW (Rq a) (Rq b) = Crypto.constEqW
        (BlockN.unsafeCast a :: SecureBlock Pub Word)
        (BlockN.unsafeCast b :: SecureBlock Pub Word)

#ifdef ML_DSA_TESTING
fromCoeffs :: [Zq] -> Maybe (Rq Sec)
fromCoeffs = fmap Rq . BlockN.fromList

toCoeffs :: Rq Sec -> [Zq]
toCoeffs (Rq a) = BlockN.toList a
#endif

newtype Tq marking = Tq (BlockN marking N Mq)
#ifdef ML_DSA_TESTING
    deriving (Eq, Show, NFData)
#else
    deriving NFData
#endif

instance Classified marking => Add (Tq marking) where
    zero = Tq zero
    Tq a .+ Tq b = Tq (a .+ b)
    {-# INLINE (.+) #-}
    Tq a .- Tq b = Tq (a .- b)
    {-# INLINE (.-) #-}
    neg (Tq a) = Tq (neg a)
    {-# INLINE neg #-}

instance Crypto.ConstEqW (Tq Pub) where
    constEqW (Tq a) (Tq b) = Crypto.constEqW
        (BlockN.unsafeCast a :: SecureBlock Pub Word)
        (BlockN.unsafeCast b :: SecureBlock Pub Word)

instance BiMul (Tq marking) (Tq Sec) where
    Tq a ..* Tq b = Tq $ BlockN.zipWithEqPrimSizeR (.*) a b
    {-# INLINE (..*) #-}

instance BiMulAdd (Tq marking) (Tq Sec) where
    biMulFold = multiplyNTTsFold
    {-# INLINE biMulFold #-}

#ifdef ML_DSA_TESTING
instance Mul (Tq Sec) where
    one = Tq $ BlockN.replicate one
    (.*) = (..*)

instance MulAdd (Tq Sec) where
    mulAdd = biMulAdd
#endif

newtype Norm = Norm { getNorm :: Word32 }
#ifdef ML_DSA_TESTING
    deriving (Eq, Ord, Show, Num)
#endif

instance Semigroup Norm where
    Norm a <> Norm b = Norm (max32 a b)

instance Monoid Norm where
    mempty = Norm 0

#ifdef ML_DSA_TESTING
norm :: Rq marking -> Norm
norm = normFrom mempty
#endif

normFrom :: Norm -> Rq marking -> Norm
normFrom (Norm m0) (Rq w) = Norm $ BlockN.foldl' f m0 w
  where
    absZq (Zq z) = min32 z (q32 - z)
    f m = max32 m . absZq
{-# INLINE normFrom #-}

countFrom :: Word -> Hints -> Word
countFrom b (Hints w) = let f acc (H h) = acc + fromIntegral h in BlockN.foldl' f b w
{-# INLINE countFrom #-}

-- Computes the NTT representation of the given polynomial
ntt :: Classified marking => Rq marking -> Tq marking
ntt (Rq a) = Tq $ BlockN.runThaw b mutNtt
    where b = BlockN.mapEqPrimSize toMontgomery a
{-# INLINE ntt #-}

mutNtt :: MutableBlockN marking N Mq s -> ST s ()
mutNtt !b = outer 1 128
  where
    outer !m len = when (len >= 1) $ inner m len 0

    inner !m !len start
        | start < 256 = do
            let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m
            loop z (start + len) len start
            inner (m + 1) len (start + offsetShiftL 1 len)
        | otherwise = outer m (offsetShiftR 1 len)

    loop !z end len j =
        when (j < end) $ do
            t <- (z .*) <$> BlockN.read b (j + len)
            x <- BlockN.read b j
            BlockN.write b (j + len) (x .- t)
            BlockN.write b j (x .+ t)
            loop z end len (j + 1)
{-# NOINLINE mutNtt #-}

-- Computes the polynomial that corresponds to the given NTT representation
nttInv :: Classified marking => Tq marking -> Rq marking
nttInv (Tq a) = Rq $ BlockN.mapEqPrimSize fromMontgomery $ BlockN.runThaw a mutNttInv
{-# INLINE nttInv #-}

mutNttInv :: MutableBlockN marking N Mq s -> ST s ()
mutNttInv !b = do
    outer 255 1
    BlockN.iterModify (\x -> x .* Mq 16382) b  -- toMontgomery (Zq 8347681)
  where
    outer !m len = when (len < 256) $ inner m len 0

    inner !m !len start
        | start < 256 = do
            let z = BlockN.index zetaPowBitRev m -- 1753 ^ bitRev8 m
            loop z (start + len) len start
            inner (m - 1) len (start + offsetShiftL 1 len)
        | otherwise = outer m (offsetShiftL 1 len)

    loop !z end len j =
        when (j < end) $ do
            t <- BlockN.read b j
            x <- BlockN.read b (j + len)
            BlockN.write b j (t .+ x)
            BlockN.write b (j + len) (z .* (x .- t))
            loop z end len (j + 1)
{-# NOINLINE mutNttInv #-}

multiplyNTTsFold :: Foldable t => Tq Sec -> t (Tq marking, Tq Sec) -> Tq Sec
multiplyNTTsFold (Tq c) = Tq . BlockN.runFold c (uncurry multiplyNTTsAdd)
{-# INLINE multiplyNTTsFold #-}

-- Multiply then add a third term
multiplyNTTsAdd :: Tq marking -> Tq Sec -> MutableBlockN Sec N Mq s -> ST s ()
multiplyNTTsAdd (Tq !f) (Tq !g) bb = loop bb 0
  where
    loop :: MutableBlockN Sec N Mq s -> Offset Mq -> ST s ()
    loop !b i = when (i < 256) $ do
        c <- BlockN.read b i
        BlockN.write b i $ mulAdd (BlockN.index f i) (BlockN.index g i) c
        loop b (i + 1)

-- Values of 1753 ^ BitRev8(𝑖) mod 𝑞 for 𝑖 ∈ {0, … , 255}
zetaPowBitRev :: BlockN Pub 256 Mq
zetaPowBitRev = BlockN.runNew (Proxy :: Proxy Pub) $ \out ->
    foldM_ (loop out) one offsets
  where
    offsets = Prelude.map (fromIntegral . bitRev8) [0 .. 255]
    loop b acc i = BlockN.write b i acc >> return (Mq 2091667 .* acc)
    -- toMontgomery (Zq 1753)

newtype H = H Word8
#ifdef ML_DSA_TESTING
    deriving (Eq, Show)
#else
    deriving Eq
#endif

instance Prim H where
    sizeOf# (H a) = sizeOf# a
    {-# INLINE sizeOf# #-}
    alignment# (H a) = alignment# a
    {-# INLINE alignment# #-}
#if MIN_VERSION_primitive(0,9,0)
    sizeOfType# _ = sizeOfType# (Proxy :: Proxy Word8)
    {-# INLINE sizeOfType# #-}
    alignmentOfType# _ = alignmentOfType# (Proxy :: Proxy Word8)
    {-# INLINE alignmentOfType# #-}
#endif
    indexByteArray# ba i = H (indexByteArray# ba i)
    {-# INLINE indexByteArray# #-}
    readByteArray# mba i s =
        case readByteArray# mba i s of
            (# s', a #) -> (# s', H a #)
    {-# INLINE readByteArray# #-}
    writeByteArray# mba i (H a) = writeByteArray# mba i a
    {-# INLINE writeByteArray# #-}
    setByteArray# mba i len (H a) = setByteArray# mba i len a
    {-# INLINE setByteArray# #-}
    indexOffAddr# addr i = H (indexOffAddr# addr i)
    {-# INLINE indexOffAddr# #-}
    readOffAddr# addr i s =
        case readOffAddr# addr i s of
            (# s', a #) -> (# s', H a #)
    {-# INLINE readOffAddr# #-}
    writeOffAddr# addr i (H a) = writeOffAddr# addr i a
    {-# INLINE writeOffAddr# #-}
    setOffAddr# addr i len (H a) = setOffAddr# addr i len a
    {-# INLINE setOffAddr# #-}

instance PrimSized H where
    type PrimSize H = 1

-- boolean-less comparison, to avoid code duplication with pattern matching
notEqualH :: Word32 -> Word32 -> H
notEqualH (W32# x#) (W32# y#) = H $
#if MIN_VERSION_base(4,16,0)
    1 - W8# (Exts.wordToWord8# (Exts.int2Word# (Exts.eqWord32# x# y#)))
#else
    1 - W8# (Exts.int2Word# (Exts.eqWord# x# y#))
#endif

newtype Hints = Hints (BlockN Sec N H)
#ifdef ML_DSA_TESTING
    deriving (Eq, Show, NFData)
#else
    deriving NFData
#endif

instance Crypto.ConstEqW Hints where
    constEqW (Hints a) (Hints b) = Crypto.constEqW
        (BlockN.unsafeCast a :: SecureBlock Sec Word)
        (BlockN.unsafeCast b :: SecureBlock Sec Word)

#ifdef ML_DSA_TESTING
fromBools :: [Bool] -> Maybe Hints
fromBools = fmap Hints . BlockN.fromList . map (H . fromIntegral . fromEnum)

toBools :: Hints -> [Bool]
toBools (Hints a) = map (\(H h) -> toEnum (fromIntegral h)) (BlockN.toList a)
#endif

peekWordPos :: Ptr WordLE -> BitPos -> ST s WordM
peekWordPos a bp = fromLE <$> ST.peekElemOff a (wordOff bp)

pokeWordPos :: Ptr WordLE -> BitPos -> WordM -> ST s ()
pokeWordPos a bp = ST.pokeElemOff a (wordOff bp) . toLE

newtype BitPos = BitPos Int

zeroPos :: BitPos
zeroPos = BitPos 0

wordOff :: BitPos -> Int
wordOff (BitPos p) = div p wordBits

bitPos :: BitPos -> Int
bitPos (BitPos p) = p .&. (wordBits - 1)

availPos :: Int -> BitPos -> Int
availPos requested (BitPos p) = min available requested
  where available = wordBits - (p .&. (wordBits - 1))

nextPos :: Int -> BitPos -> (Int, BitPos)
nextPos requested (BitPos p) = (howMany, BitPos $ p + howMany)
  where howMany = availPos requested (BitPos p)

getMask :: Int -> WordM
getMask howMany
    | howMany >= wordBits = maxBound
    | otherwise = (1 `unsafeShiftL` howMany) - 1
    -- branch useful only when processing one byte at a time due to
    -- architecture not supporting unaligned memory access

-- Encodes an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 22
simpleBitPack :: Int -> BlockN marking N Word32 -> Builder marking
simpleBitPack d w = Builder.create (32 * d) (runSimpleBitPack d w)
{-# INLINE simpleBitPack #-}

runSimpleBitPack :: Int -> BlockN marking N Word32 -> Ptr WordLE -> ST s ()
runSimpleBitPack !d !w dst = loop dst 0 zeroPos 0 (get 0) d
  where
    get = BlockN.index w . Offset
    {-# INLINE get #-}

    loop !b !pos !bp !o !a j
        | j == 0, pos' == n = return ()
        | j == 0 = loop b pos' bp o (get pos') d
        | bitPos bp + howMany < wordBits = loop b pos bp' o' a' j'
        | otherwise = pokeWordPos b bp o' >> loop b pos bp' 0 a' j'
      where
        pos' = pos + 1
        (howMany, bp') = nextPos j bp
        x = fromIntegral a .&. getMask howMany
        o' = o .|. (x `unsafeShiftL` bitPos bp)
        a' = a `unsafeShiftR` howMany
        j' = j - howMany

-- simpleBitPack with 𝑑=10 after conversion from a polynomial and division by 2^13
simpleBitPack10 :: Rq Pub -> Builder Pub
simpleBitPack10 (Rq w) = simpleBitPack 10 (BlockN.mapEqPrimSize f w)
  where f (Zq z) = z `unsafeShiftR` 13
{-# INLINE simpleBitPack10 #-}

-- Decodes an array of 𝑑-bit integers for 1 ≤ 𝑑 ≤ 22
simpleBitUnpack :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> BlockN marking N Word32
simpleBitUnpack d b = BlockN.runNew Proxy $ mutSimpleBitUnpack d b
{-# INLINE simpleBitUnpack #-}

mutSimpleBitUnpack :: ByteArrayAccess ba => Int -> ba -> MutableBlockN marking N Word32 s -> ST s ()
mutSimpleBitUnpack !d b !w = ST.withByteArray b $ \p -> outer p zeroPos 0
  where
    outer !p !bp i = when (i < Offset n) $ inner p i bp 0 0

    inner !p !i !bp !v j
        | j == d = BlockN.write w i v >> outer p bp (i + 1)
        | otherwise = do
            let (howMany, bp') = nextPos (d - j) bp
            y <- get p bp howMany
            let v' = v .|. (fromIntegral y `unsafeShiftL` j)
                j' = j + howMany
            inner p i bp' v' j'

    get :: Ptr WordLE -> BitPos -> Int -> ST s WordM
    get p bp howMany = do
        x <- (`unsafeShiftR` bitPos bp) <$> peekWordPos p bp
        return (x .&. getMask howMany)

-- simpleBitUnpack with 𝑑=10 and conversion to a polynomial with multiplication by 2^13
simpleBitUnpack10 :: ByteArrayAccess ba => ba -> Rq Pub
simpleBitUnpack10 = Rq . BlockN.mapEqPrimSize f . simpleBitUnpack 10
  where f z = Zq (z `unsafeShiftL` 13)
{-# INLINE simpleBitUnpack10 #-}

-- bitPack when upper bound is a power of two
bitPackSafe :: Classified marking => Int -> Rq marking -> Builder marking
bitPackSafe d = simpleBitPack d . unmirror (1 `unsafeShiftL` (d - 1))
{-# INLINE bitPackSafe #-}

-- bitUnpack when upper bound is a power of two
bitUnpackSafe :: (Classified marking, ByteArrayAccess ba) => Int -> ba -> Rq marking
bitUnpackSafe d b = mirror (1 `unsafeShiftL` (d - 1)) (simpleBitUnpack d b)
{-# INLINE bitUnpackSafe #-}

-- Encodes a polynomial 𝑤 into a byte string
bitPack :: Word32 -> Int -> Rq Sec -> Builder Sec
bitPack !m d = simpleBitPack d . unmirror m

-- Reverses the procedure bitPack
bitUnpack :: ByteArrayAccess ba => Word32 -> Int -> ba -> Maybe (Rq Sec)
bitUnpack !m d b
    | valid unpacked = Just $ mirror m unpacked
    | otherwise = Nothing
  where
    unpacked  = simpleBitUnpack d b
    valid     = Crypto.toBool . BlockN.foldl' f Crypto.trueW
    f acc z   = acc `Crypto.andW` inRange z
    inRange z = Crypto.ltW (fromIntegral z) (fromIntegral (2 * m + 1))

mirror :: Classified marking => Word32 -> BlockN marking N Word32 -> Rq marking
mirror m w = Rq $ BlockN.mapEqPrimSize (\z -> Zq m .- Zq z) w
{-# INLINE mirror #-}

unmirror :: Classified marking => Word32 -> Rq marking -> BlockN marking N Word32
unmirror m (Rq w) = BlockN.mapEqPrimSize (\z -> fromZq (Zq m .- z)) w
{-# INLINE unmirror #-}

-- Encodes a polynomial vector 𝐡 with binary coefficients into a byte string
hintBitPack :: forall k. KnownNat k => Int -> Vector k Hints -> Builder Sec
hintBitPack !omega !h = Builder.create (omega + k) $ outer 0 0
  where
    k = fromIntegral $ natVal (Proxy :: Proxy k)

    outer :: Int -> Offset Word8 -> Ptr Word8 -> ST s ()
    outer i !idx !p
        | i == k = end p idx
        | otherwise = inner p i (Vector.index h $ Offset i) idx 0

    inner :: Ptr Word8 -> Int -> Hints -> Offset Word8 -> Int -> ST s ()
    inner !p !i e@(Hints !w) !idx j
        | j == n = mark p i idx >> outer (i + 1) idx p
        | BlockN.index w (Offset j) == H 0 = inner p i e idx (j + 1)
        | otherwise = set p idx j >> inner p i e (idx + 1) (j + 1)

    set :: Ptr Word8 -> Offset Word8 -> Int -> ST s ()
    set p (Offset idx) j = ST.pokeElemOff p idx (fromIntegral j)

    mark :: Ptr Word8 -> Int -> Offset Word8 -> ST s ()
    mark p i (Offset idx) = ST.pokeElemOff p (omega + i) (fromIntegral idx)

    end :: Ptr Word8 -> Offset Word8 -> ST s ()
    end !p (Offset idx) =
        when (idx < omega) $ ST.fillBytes (p `plusPtr` idx) 0 (omega - idx)
{-# INLINE hintBitPack #-}

-- Reverses the procedure hintBitPack
hintBitUnpack :: (KnownNat k, ByteArrayAccess ba) => Int -> ba -> Maybe (Vector k Hints)
hintBitUnpack !omega !y = Vector.createMaybeAccum outer cont 0
  where
    outer :: Int -> Offset Hints -> Maybe (Int, Hints)
    outer idx (Offset i) = runST $ ST.withByteArray y $ \p -> do
        m <- fromIntegral <$> ST.peekElemOff p (omega + i)
        if m < idx || m > omega
            then return Nothing
            else new >>= inner p m idx

    inner :: Ptr Word8 -> Int -> Int -> MutableBlockN Sec N H s -> ST s (Maybe (Int, Hints))
    inner !p !m !first !mb = go first
      where
        go idx
            | idx >= m = unsafeFreezeF mb >>= \b ->
                return $ Just (idx, Hints b)
            | idx > first = do
                val  <- ST.peekElemOff p idx
                val' <- ST.peekElemOff p (idx - 1)
                if val' >= val
                    then return Nothing
                    else set mb val >> go (idx + 1)
            | otherwise = ST.peekElemOff p idx >>= \val ->
                set mb val >> go (idx + 1)

    new :: ST s (MutableBlockN Sec N H s)
    new = newF >>= \mb -> BlockN.erase mb >> return mb

    set :: MutableBlockN Sec N H s -> Word8 -> ST s ()
    set mb val = BlockN.write mb (fromIntegral val) (H 1)

    cont :: forall r. Int -> r -> Maybe r
    cont idx h = runST $ ST.withByteArray y $ \p -> go p idx
      where
        go :: Ptr Word8 -> Int -> ST s (Maybe r)
        go !p i
            | i >= omega = return (Just h)
            | otherwise = do
                val <- ST.peekElemOff p i
                case val of
                    0 -> go p (i + 1)
                    _ -> return Nothing
{-# INLINE hintBitUnpack #-}

-- Samples a polynomial 𝑐 ∈ 𝑅 with coefficients from {-1, 0, 1} and
-- Hamming weight 𝜏 ≤ 64
sampleInBall :: Int -> SecureBytes Sec -> Rq Sec
sampleInBall tau rho = Rq $
    BlockN.runNew (Proxy :: Proxy Sec) $ \c -> runXof c 136
  where
    runXof :: MutableBlockN Sec N Zq s -> Int -> ST s ()
    runXof !c !xofLen = case someNatVal (fromIntegral (8 * xofLen)) of
        SomeNat proxy -> do
            let bytes = Crypto.unDigest (doHash proxy)
            BlockN.erase c
            loop c xofLen bytes 8 (256 - tau)

    loop :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Int -> ST s ()
    loop !c !xofLen !bytes !pos i
        | i == n = return ()
        | pos >= xofLen = runXof c (xofLen + 136)
        | j > Offset i = loop c xofLen bytes (pos + 1) i
        | otherwise = do
            BlockN.read c j >>= BlockN.write c (Offset i)
            let idx   = i + tau - 256
                iByte = idx `unsafeShiftR` 3
                iBit  = idx .&. 7
                b = B.index bytes iByte `unsafeShiftR` iBit
            BlockN.write c j (Zq $ if even b then 1 else q32 - 1)
            loop c xofLen bytes (pos + 1) (i + 1)
      where
        j = fromIntegral $ B.index bytes pos

    doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)
    doHash _ = Crypto.hash rho

-- Samples a polynomial ∈ 𝑇𝑞
rejNttPoly :: SecureBytes Pub -> Word8 -> Word8 -> Tq Pub
rejNttPoly rho !s !r = Tq $
    BlockN.runNew (Proxy :: Proxy Pub) $ \b -> runXof b (280 * 3) 0 0
  where
    runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of
        SomeNat proxy -> do
            let bytes = Crypto.unBlockDigest (doHash proxy)
            loop b xofLen bytes pos j

    loop !b !xofLen !bytes !pos j
        | j == 256 = return ()
        | pos >= Offset xofLen = runXof b (xofLen + 56 * 3) pos j
        | otherwise = do
            let s0 = fromIntegral $ blockIndex bytes pos
                s1 = fromIntegral $ blockIndex bytes (pos + 1)
                s2 = fromIntegral $ blockIndex bytes (pos + 2)
                z = s0 + (s1 `unsafeShiftL` 8)
                       + ((s2 .&. 0x7F) `unsafeShiftL` 16)
            poke b j z >>= loop b xofLen bytes (pos + 3)

    poke b j z
        | z < q32 = BlockN.write b j (toMontgomery $ Zq z) >> return (j + 1)
        | otherwise = return j

    doHash :: KnownNat bitlen => proxy bitlen -> BlockDigest (SHAKE128 bitlen)
    doHash _ = Crypto.hashToBlock input

    input :: SecureBytes Pub
    !input = B.unsafeCreate (len + 2) $ \d -> do
        B.copyByteArrayToPtr rho d
        pokeByteOff d len s
        pokeByteOff d (len + 1) r
    len = B.length rho

-- Samples an element 𝑎 ∈ 𝑅 with coefficients in [-𝜂, 𝜂] computed via rejection
-- sampling from 𝜌
rejBoundedPoly :: SecureBytes Sec -> Int -> Word16 -> Rq Sec
rejBoundedPoly rho !eta !r = Rq $
    BlockN.runNew (Proxy :: Proxy Sec) $ \b -> runXof b 408 0 0
  where
    runXof :: MutableBlockN Sec N Zq s -> Int -> Int -> Offset Zq -> ST s ()
    runXof !b !xofLen !pos !j = case someNatVal (fromIntegral (8 * xofLen)) of
        SomeNat proxy -> case eta of
            2 -> loop2 b xofLen bytes pos j
            4 -> loop4 b xofLen bytes pos j
            _ -> error "rejBoundedPoly: unexpected eta value"
          where bytes = Crypto.unDigest (doHash proxy)

    loop2 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s ()
    loop2 !b !xofLen !bytes !pos j
        | j == 256 = return ()
        | pos >= xofLen = runXof b (xofLen + 136) pos j
        | otherwise = do
            let z  = fromIntegral $ B.index bytes pos
                z0 = z .&. 0xF
                z1 = z `unsafeShiftR` 4
            j2 <- poke2 b j z0
            when (j2 < 256) $ poke2 b j2 z1 >>= loop2 b xofLen bytes (pos + 1)

    loop4 :: MutableBlockN Sec N Zq s -> Int -> SecureBytes Sec -> Int -> Offset Zq -> ST s ()
    loop4 !b !xofLen !bytes !pos j
        | j == 256 = return ()
        | pos >= xofLen = runXof b (xofLen + 136) pos j
        | otherwise = do
            let z  = fromIntegral $ B.index bytes pos
                z0 = z .&. 0xF
                z1 = z `unsafeShiftR` 4
            j2 <- poke4 b j z0
            when (j2 < 256) $ poke4 b j2 z1 >>= loop4 b xofLen bytes (pos + 1)

    -- rejection sampling from {-2, … , 2}
    poke2 b j z
        | z < 15 = BlockN.write b j (Zq 2 .- Zq (mod5 z)) >> return (j + 1)
        | otherwise = return j

    -- rejection sampling from {-4, … , 4}
    poke4 b j z
        | z < 9 = BlockN.write b j (Zq 4 .- Zq z) >> return (j + 1)
        | otherwise = return j

    doHash :: KnownNat bitlen => proxy bitlen -> Digest (SHAKE256 bitlen)
    doHash _ = Crypto.hash input

    input :: SecureBytes Sec
    !input = B.unsafeCreate (len + 2) $ \d -> do
        B.copyByteArrayToPtr rho d
        pokeByteOff d len (B.toLE r)
    len = B.length rho

powerTwoRound :: Rq Sec -> (Rq Pub, Rq Pub)
powerTwoRound (Rq !a) = runST $ do
    mb <- newF
    mc <- newF
    loop mb mc 0
    b <- unsafeFreezeF mb
    c <- unsafeFreezeF mc
    return (Rq b, Rq c)
  where
    loop :: MutableBlockN Pub N Zq s -> MutableBlockN Pub N Zq s -> Int -> ST s ()
    loop !mb !mc i = when (i < n) $ do
        let r = BlockN.index a (Offset i)
            (r1, r0) = powerTwoRoundZq r
        BlockN.write mb (Offset i) r1
        BlockN.write mc (Offset i) r0
        loop mb mc (i + 1)

-- Decomposes 𝑟 into (𝑟1, 𝑟0) such that 𝑟 ≡ 𝑟1.2^𝑑 + 𝑟0 mod 𝑞 (for 𝑑=13)
powerTwoRoundZq :: Zq -> (Zq, Zq)
powerTwoRoundZq (Zq r) = (Zq r1, Zq r0)
  where
    r0' = r .&. 0x1fff
    r1' = r .&. 0xffffe000 -- keep multiplied by 2^13

    -- if r0 > 4096, increment r1 and subtract 8192 from r0
    cond = (4096 - r0') `unsafeShiftIR` 31
    r0   = select32 cond (q32 - 0x2000 + r0') r0'
    r1   = reduceSimple (r1' + (cond .&. 0x2000))

-- Decomposes 𝑟 into (𝑟1, 𝑟0) such that 𝑟 ≡ 𝑟1.(2𝛾2) + 𝑟0 mod 𝑞
decomposeZq :: Word32 -> Zq -> (Word32, Zq)
decomposeZq gamma2 r@(Zq rr) = (r1, Zq r0)
  where
    r1  = highZq gamma2 r
    r0' = rr - r1 * (2 * gamma2)
    !r0 = r0' + ((r0' `unsafeShiftIR` 31) .&. q32)

highZq :: Word32 -> Zq -> Word32
highZq gamma2 (Zq r) = case gamma2 of
    95232 -> high190464 r
    _     -> high523776 r

high190464 :: Word32 -> Word32
high190464 r = r1 `xor` (((43 - r1) `unsafeShiftIR` 31) .&. r1)
  where
    r' = (r + 127) `unsafeShiftR` 7
    r1 = (r' * 11275 + (1 `unsafeShiftL` 23)) `unsafeShiftR` 24;

high523776 :: Word32 -> Word32
high523776 r = r1 .&. 15
  where
    r' = (r + 127) `unsafeShiftR` 7
    r1 = (r' * 1025 + (1 `unsafeShiftL` 21)) `unsafeShiftR` 22;

-- Returns 𝑟1 from the output of decomposeZq 𝑟
highBits :: Word32 -> Rq Sec -> BlockN Sec N Word32
highBits gamma2 (Rq w) = BlockN.seq gamma2 $
    BlockN.mapEqPrimSize (highZq gamma2) w
{-# INLINE highBits #-}

-- Returns 𝑟0 from the output of decomposeZq 𝑟
lowBits :: Word32 -> Rq Sec -> Rq Sec
lowBits gamma2 (Rq w) = Rq $ BlockN.seq gamma2 $
    BlockN.mapEqPrimSize (snd . decomposeZq gamma2) w
{-# INLINE lowBits #-}

-- Computes hint bit indicating whether adding 𝑧 to 𝑟 alters the high bits of 𝑟
makeHint :: Word32 -> Rq Sec -> Rq Sec -> Hints
makeHint gamma2 (Rq z) (Rq r) = Hints $ BlockN.seq gamma2 $
    BlockN.zipWith (makeHintH gamma2) z r
{-# INLINE makeHint #-}

-- Returns the high bits of 𝑟 adjusted according to hint ℎ
useHint :: Word32 -> Hints -> Rq Sec -> BlockN Sec N Word32
useHint gamma2 (Hints h) (Rq r) = BlockN.seq gamma2 $
    BlockN.zipWithEqPrimSizeR (useHintH gamma2) h r
{-# INLINE useHint #-}

makeHintH :: Word32 -> Zq -> Zq -> H
makeHintH gamma2 z r = notEqualH r1 v1
  where
    r1 = highZq gamma2 r
    v1 = highZq gamma2 (r .+ z)

useHintH :: Word32 -> H -> Zq -> Word32
useHintH !gamma2 h r
    | h == H 0    = r1
    | positive r0 = decr gamma2 r1
    | otherwise   = incr gamma2 r1
  where
    (r1, Zq r0) = decomposeZq gamma2 r

    -- useHintH already has a 6-way code path with just 3 possibilities for
    -- output times the match on gamma2, so for the other tests we do not use
    -- boolean expressions, otherwise code size would grow even more

    -- z > 0 and z <= (𝑞-1)/2
    positive z =
        let x = negate z
            m = (q32 - 1) `div` 2
            w = x .&. xor (m + x) 0x80000000
         in w < 0x80000000

    incr43, decr43 :: Word32 -> Word32
    decr43 x = 43 - incr43 (43 - x)
    incr43 x = (x + 1) .&. complement wrapAroundMask
      where wrapAroundMask = (42 - x) `unsafeShiftIR` 31

    incr 95232 z = incr43 z         -- if z == 43 then 0 else z + 1
    incr _     z = (z + 1) .&. 15

    decr 95232 z = decr43 z         -- if z == 0 then 43 else z - 1
    decr _     z = (z - 1) .&. 15
{-# INLINE useHintH #-}