-- |
-- 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 DataKinds #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
module Base
    ( CountOf(..), KnownNat, Nat, Offset(..), PrimMonad, PrimType, PrimSized(..)
    , PrimState, natVal, offsetShiftL, offsetShiftR, unsafePrimFromIO, (.==#)
#ifdef ML_KEM_TESTING
    , checkBounds
#endif
    ) where

import Control.Monad.Primitive

import Data.Primitive.Types (Prim)

import Data.Bits
import Data.Word

import GHC.TypeNats

type PrimType = Prim

class PrimType a => PrimSized a where
    type PrimSize a :: Nat

instance PrimSized Word16 where
    type PrimSize Word16 = 2

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

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
