-- |
-- 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 (Int -> CountOf ty -> ShowS
[CountOf ty] -> ShowS
CountOf ty -> String
(Int -> CountOf ty -> ShowS)
-> (CountOf ty -> String)
-> ([CountOf ty] -> ShowS)
-> Show (CountOf ty)
forall ty. Int -> CountOf ty -> ShowS
forall ty. [CountOf ty] -> ShowS
forall ty. CountOf ty -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall ty. Int -> CountOf ty -> ShowS
showsPrec :: Int -> CountOf ty -> ShowS
$cshow :: forall ty. CountOf ty -> String
show :: CountOf ty -> String
$cshowList :: forall ty. [CountOf ty] -> ShowS
showList :: [CountOf ty] -> ShowS
Show, CountOf ty -> CountOf ty -> Bool
(CountOf ty -> CountOf ty -> Bool)
-> (CountOf ty -> CountOf ty -> Bool) -> Eq (CountOf ty)
forall ty. CountOf ty -> CountOf ty -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall ty. CountOf ty -> CountOf ty -> Bool
== :: CountOf ty -> CountOf ty -> Bool
$c/= :: forall ty. CountOf ty -> CountOf ty -> Bool
/= :: CountOf ty -> CountOf ty -> Bool
Eq, Eq (CountOf ty)
Eq (CountOf ty) =>
(CountOf ty -> CountOf ty -> Ordering)
-> (CountOf ty -> CountOf ty -> Bool)
-> (CountOf ty -> CountOf ty -> Bool)
-> (CountOf ty -> CountOf ty -> Bool)
-> (CountOf ty -> CountOf ty -> Bool)
-> (CountOf ty -> CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty -> CountOf ty)
-> Ord (CountOf ty)
CountOf ty -> CountOf ty -> Bool
CountOf ty -> CountOf ty -> Ordering
CountOf ty -> CountOf ty -> CountOf ty
forall ty. Eq (CountOf ty)
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall ty. CountOf ty -> CountOf ty -> Bool
forall ty. CountOf ty -> CountOf ty -> Ordering
forall ty. CountOf ty -> CountOf ty -> CountOf ty
$ccompare :: forall ty. CountOf ty -> CountOf ty -> Ordering
compare :: CountOf ty -> CountOf ty -> Ordering
$c< :: forall ty. CountOf ty -> CountOf ty -> Bool
< :: CountOf ty -> CountOf ty -> Bool
$c<= :: forall ty. CountOf ty -> CountOf ty -> Bool
<= :: CountOf ty -> CountOf ty -> Bool
$c> :: forall ty. CountOf ty -> CountOf ty -> Bool
> :: CountOf ty -> CountOf ty -> Bool
$c>= :: forall ty. CountOf ty -> CountOf ty -> Bool
>= :: CountOf ty -> CountOf ty -> Bool
$cmax :: forall ty. CountOf ty -> CountOf ty -> CountOf ty
max :: CountOf ty -> CountOf ty -> CountOf ty
$cmin :: forall ty. CountOf ty -> CountOf ty -> CountOf ty
min :: CountOf ty -> CountOf ty -> CountOf ty
Ord, Integer -> CountOf ty
CountOf ty -> CountOf ty
CountOf ty -> CountOf ty -> CountOf ty
(CountOf ty -> CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty)
-> (CountOf ty -> CountOf ty)
-> (Integer -> CountOf ty)
-> Num (CountOf ty)
forall ty. Integer -> CountOf ty
forall ty. CountOf ty -> CountOf ty
forall ty. CountOf ty -> CountOf ty -> CountOf ty
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: forall ty. CountOf ty -> CountOf ty -> CountOf ty
+ :: CountOf ty -> CountOf ty -> CountOf ty
$c- :: forall ty. CountOf ty -> CountOf ty -> CountOf ty
- :: CountOf ty -> CountOf ty -> CountOf ty
$c* :: forall ty. CountOf ty -> CountOf ty -> CountOf ty
* :: CountOf ty -> CountOf ty -> CountOf ty
$cnegate :: forall ty. CountOf ty -> CountOf ty
negate :: CountOf ty -> CountOf ty
$cabs :: forall ty. CountOf ty -> CountOf ty
abs :: CountOf ty -> CountOf ty
$csignum :: forall ty. CountOf ty -> CountOf ty
signum :: CountOf ty -> CountOf ty
$cfromInteger :: forall ty. Integer -> CountOf ty
fromInteger :: Integer -> CountOf ty
Num)

newtype Offset ty = Offset Int
    deriving (Int -> Offset ty -> ShowS
[Offset ty] -> ShowS
Offset ty -> String
(Int -> Offset ty -> ShowS)
-> (Offset ty -> String)
-> ([Offset ty] -> ShowS)
-> Show (Offset ty)
forall ty. Int -> Offset ty -> ShowS
forall ty. [Offset ty] -> ShowS
forall ty. Offset ty -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall ty. Int -> Offset ty -> ShowS
showsPrec :: Int -> Offset ty -> ShowS
$cshow :: forall ty. Offset ty -> String
show :: Offset ty -> String
$cshowList :: forall ty. [Offset ty] -> ShowS
showList :: [Offset ty] -> ShowS
Show, Offset ty -> Offset ty -> Bool
(Offset ty -> Offset ty -> Bool)
-> (Offset ty -> Offset ty -> Bool) -> Eq (Offset ty)
forall ty. Offset ty -> Offset ty -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall ty. Offset ty -> Offset ty -> Bool
== :: Offset ty -> Offset ty -> Bool
$c/= :: forall ty. Offset ty -> Offset ty -> Bool
/= :: Offset ty -> Offset ty -> Bool
Eq, Eq (Offset ty)
Eq (Offset ty) =>
(Offset ty -> Offset ty -> Ordering)
-> (Offset ty -> Offset ty -> Bool)
-> (Offset ty -> Offset ty -> Bool)
-> (Offset ty -> Offset ty -> Bool)
-> (Offset ty -> Offset ty -> Bool)
-> (Offset ty -> Offset ty -> Offset ty)
-> (Offset ty -> Offset ty -> Offset ty)
-> Ord (Offset ty)
Offset ty -> Offset ty -> Bool
Offset ty -> Offset ty -> Ordering
Offset ty -> Offset ty -> Offset ty
forall ty. Eq (Offset ty)
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall ty. Offset ty -> Offset ty -> Bool
forall ty. Offset ty -> Offset ty -> Ordering
forall ty. Offset ty -> Offset ty -> Offset ty
$ccompare :: forall ty. Offset ty -> Offset ty -> Ordering
compare :: Offset ty -> Offset ty -> Ordering
$c< :: forall ty. Offset ty -> Offset ty -> Bool
< :: Offset ty -> Offset ty -> Bool
$c<= :: forall ty. Offset ty -> Offset ty -> Bool
<= :: Offset ty -> Offset ty -> Bool
$c> :: forall ty. Offset ty -> Offset ty -> Bool
> :: Offset ty -> Offset ty -> Bool
$c>= :: forall ty. Offset ty -> Offset ty -> Bool
>= :: Offset ty -> Offset ty -> Bool
$cmax :: forall ty. Offset ty -> Offset ty -> Offset ty
max :: Offset ty -> Offset ty -> Offset ty
$cmin :: forall ty. Offset ty -> Offset ty -> Offset ty
min :: Offset ty -> Offset ty -> Offset ty
Ord, Integer -> Offset ty
Offset ty -> Offset ty
Offset ty -> Offset ty -> Offset ty
(Offset ty -> Offset ty -> Offset ty)
-> (Offset ty -> Offset ty -> Offset ty)
-> (Offset ty -> Offset ty -> Offset ty)
-> (Offset ty -> Offset ty)
-> (Offset ty -> Offset ty)
-> (Offset ty -> Offset ty)
-> (Integer -> Offset ty)
-> Num (Offset ty)
forall ty. Integer -> Offset ty
forall ty. Offset ty -> Offset ty
forall ty. Offset ty -> Offset ty -> Offset ty
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: forall ty. Offset ty -> Offset ty -> Offset ty
+ :: Offset ty -> Offset ty -> Offset ty
$c- :: forall ty. Offset ty -> Offset ty -> Offset ty
- :: Offset ty -> Offset ty -> Offset ty
$c* :: forall ty. Offset ty -> Offset ty -> Offset ty
* :: Offset ty -> Offset ty -> Offset ty
$cnegate :: forall ty. Offset ty -> Offset ty
negate :: Offset ty -> Offset ty
$cabs :: forall ty. Offset ty -> Offset ty
abs :: Offset ty -> Offset ty
$csignum :: forall ty. Offset ty -> Offset ty
signum :: Offset ty -> Offset ty
$cfromInteger :: forall ty. Integer -> Offset ty
fromInteger :: Integer -> Offset ty
Num)

offsetShiftL :: Int -> Offset ty -> Offset ty2
offsetShiftL :: forall ty ty2. Int -> Offset ty -> Offset ty2
offsetShiftL Int
n (Offset Int
o) = Int -> Offset ty2
forall ty. Int -> Offset ty
Offset (Int
o Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
n)

offsetShiftR :: Int -> Offset ty -> Offset ty2
offsetShiftR :: forall ty ty2. Int -> Offset ty -> Offset ty2
offsetShiftR Int
n (Offset Int
o) = Int -> Offset ty2
forall ty. Int -> Offset ty
Offset (Int
o Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
n)

unsafePrimFromIO :: PrimMonad m => IO a -> m a
unsafePrimFromIO :: forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafePrimFromIO = IO a -> m a
forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim

(.==#) :: Offset ty -> CountOf ty -> Bool
.==# :: forall ty. Offset ty -> CountOf ty -> Bool
(.==#) (Offset Int
ofs) (CountOf Int
sz) = Int
ofs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
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