mlkem-0.2.0.0: src/Base.hs

-- |
-- Module      : Base
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- Utilities similar to ones provided in basement but adapted for primitive
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Base
    ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimState
    , natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#)
#ifdef ML_KEM_TESTING
    , checkBounds
#endif
    ) where

import Control.Monad.Primitive

import Data.Primitive.Types

import Data.Bits

import GHC.TypeLits

type PrimType = Prim

newtype CountOf ty = CountOf Int
    deriving (Show, Eq, Ord)

newtype Offset ty = Offset Int
    deriving (Show, Eq, Ord, Num)

offsetShiftL :: Int -> Offset ty -> Offset ty2
offsetShiftL n (Offset o) = Offset (o `unsafeShiftL` n)

offsetShiftR :: Int -> Offset ty -> Offset ty2
offsetShiftR n (Offset o) = Offset (o `unsafeShiftR` n)

unsafePrimFromIO :: PrimMonad m => IO a -> m a
unsafePrimFromIO = unsafeIOToPrim

(.==#) :: Offset ty -> CountOf ty -> Bool
(.==#) (Offset ofs) (CountOf sz) = ofs == sz
{-# INLINE (.==#) #-}

#ifdef ML_KEM_TESTING
checkBounds :: CountOf ty -> Offset ty -> a -> a
checkBounds (CountOf sz) (Offset ofs) a
    | ofs >= 0 && ofs < sz = a
    | otherwise = error ("offset not in valid range: " ++ show ofs)
#endif