mldsa-0.1.0.0: src/Builder.hs

-- |
-- Module      : Builder
-- License     : BSD-3-Clause
-- Copyright   : (c) 2025 Olivier Chéron
--
-- Eliminate intermediate allocations when concatenating several buffers
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
module Builder
    ( Builder, builderLength, bytes, copyBuilderToPtr, create, promote, public
    , run, runRelaxed, runToBlock, secret, storable, unsafeCreate
    ) where

import Data.ByteArray (ByteArray, ByteArrayAccess)
import qualified Data.ByteArray as B

import Control.Monad.ST
import Control.Monad.ST.Unsafe

import Data.Semigroup
import Data.Word

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

import qualified GHC.Exts as Exts

import Base
import Block (Block)
import Marking (Classified, Leak(..), SecurityMarking(..))
import SecureBytes (SecureBytes)
import qualified Block
import qualified ByteArrayST as ST
import qualified SecureBytes

data Builder (marking :: SecurityMarking) = Builder
    { builderLength :: {-# UNPACK #-} !Int
    , copyBuilderToPtr :: forall s. Ptr Word8 -> ST s ()
    }

instance Semigroup (Builder marking) where
    b1 <> b2  = create (n1 + n2) $ \p ->
        copyBuilderToPtr b1 p >> copyBuilderToPtr b2 (p `plusPtr` n1)
      where
        n1 = builderLength b1
        n2 = builderLength b2

instance Monoid (Builder marking) where
    mempty = empty
    mconcat builders = create (getSize (Exts.inline builders)) $
        go (Exts.inline builders)
      where
        getSize = getSum . Prelude.mconcat . map (Sum . builderLength)
        go = foldr c (const $ return ())
        c b k = Exts.oneShot $ \p -> copyBuilderToPtr b p >> k (p `plusPtr` builderLength b)
    {-# INLINE mconcat #-}

instance Leak Builder

bytes :: Classified marking => SecureBytes marking -> Builder marking
bytes b = unsafeCreate (SecureBytes.length b) (SecureBytes.copyByteArrayToPtr b)

create :: Int -> (forall s. Ptr a -> ST s ()) -> Builder marking
create n f = Builder n (f . castPtr)
{-# INLINE create #-}

empty :: Builder marking
empty = Builder 0 $ \_ -> return ()

public :: ByteArrayAccess ba => ba -> Builder marking
public b = unsafeCreate (B.length b) (B.copyByteArrayToPtr b)

promote :: Builder Pub -> Builder Sec
promote (Builder n f) = Builder n f

secret :: ByteArrayAccess ba => ba -> Builder Sec
secret = public

run :: Classified marking => Builder marking -> SecureBytes marking
run b = SecureBytes.unsafeCreate (builderLength b) (copyBuilderToPtr b)

runRelaxed :: ByteArray ba => Builder Pub -> ba
runRelaxed b = ST.unsafeCreate (builderLength b) (copyBuilderToPtr b)

runToBlock :: Builder Pub -> Block Word8
runToBlock b = runST $ do
    mb <- Block.newPinned (CountOf $ builderLength b)
    copyBuilderToPtr b (Block.mutableContents mb)
    Block.unsafeFreeze mb

storable :: Storable a => a -> Builder marking
storable !a = unsafeCreate (sizeOf a) (`poke` a)

unsafeCreate :: Int -> (Ptr a -> IO ()) -> Builder marking
unsafeCreate n f = create n (unsafeIOToST . f)
{-# INLINE unsafeCreate #-}