-- |
-- Module      : Crypto.Internal.WordArray
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : Good
--
-- Small and self contained array representation
-- with limited safety for internal use.
--
-- The array produced should never be exposed to the user directly.
--
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module Crypto.Internal.WordArray
    ( Array8
    , Array32
    , Array64
    , MutableArray32
    , array8
    , array32
    , array32FromAddrBE
    , allocArray32AndFreeze
    , mutableArray32
    , array64
    , arrayRead8
    , arrayRead32
    , arrayRead64
    , mutableArrayRead32
    , mutableArrayWrite32
    , mutableArrayWriteXor32
    , mutableArray32FromAddrBE
    , mutableArray32Freeze
    ) where

import Data.Word
import Data.Bits (xor)
import Crypto.Internal.Compat
import Crypto.Internal.CompatPrim
import GHC.Prim
import GHC.Types
import GHC.Word

-- | Array of Word8
data Array8 = Array8 Addr#

-- | Array of Word32
data Array32 = Array32 ByteArray#

-- | Array of Word64
data Array64 = Array64 ByteArray#

-- | Array of mutable Word32
data MutableArray32 = MutableArray32 (MutableByteArray# RealWorld)

-- | Create an array of Word8 aliasing an Addr#
array8 :: Addr# -> Array8
array8 :: Addr# -> Array8
array8 = Addr# -> Array8
Array8

-- | Create an Array of Word32 of specific size from a list of Word32
array32 :: Int -> [Word32] -> Array32
array32 :: Int -> [Word32] -> Array32
array32 Int
n [Word32]
l = forall a. IO a -> a
unsafeDoIO (Int -> [Word32] -> IO MutableArray32
mutableArray32 Int
n [Word32]
l forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MutableArray32 -> IO Array32
mutableArray32Freeze)
{-# NOINLINE array32 #-}

-- | Create an Array of BE Word32 aliasing an Addr
array32FromAddrBE :: Int -> Addr# -> Array32
array32FromAddrBE :: Int -> Addr# -> Array32
array32FromAddrBE Int
n Addr#
a =
    forall a. IO a -> a
unsafeDoIO (Int -> Addr# -> IO MutableArray32
mutableArray32FromAddrBE Int
n Addr#
a forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MutableArray32 -> IO Array32
mutableArray32Freeze)
{-# NOINLINE array32FromAddrBE #-}

-- | Create an Array of Word32 using an initializer
allocArray32AndFreeze :: Int -> (MutableArray32 -> IO ()) -> Array32
allocArray32AndFreeze :: Int -> (MutableArray32 -> IO ()) -> Array32
allocArray32AndFreeze Int
n MutableArray32 -> IO ()
f =
    forall a. IO a -> a
unsafeDoIO (Int -> [Word32] -> IO MutableArray32
mutableArray32 Int
n [] forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \MutableArray32
m -> MutableArray32 -> IO ()
f MutableArray32
m forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MutableArray32 -> IO Array32
mutableArray32Freeze MutableArray32
m)
{-# NOINLINE allocArray32AndFreeze #-}

-- | Create an Array of Word64 of specific size from a list of Word64
array64 :: Int -> [Word64] -> Array64
array64 :: Int -> [Word64] -> Array64
array64 (I# Int#
n) [Word64]
l = forall a. IO a -> a
unsafeDoIO forall a b. (a -> b) -> a -> b
$ forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
    case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (Int#
n Int# -> Int# -> Int#
*# Int#
8#) Int#
8# State# RealWorld
s of
        (# State# RealWorld
s', MutableByteArray# RealWorld
mbarr #) -> forall {d}.
Int#
-> State# d
-> MutableByteArray# d
-> [Word64]
-> (# State# d, Array64 #)
loop Int#
0# State# RealWorld
s' MutableByteArray# RealWorld
mbarr [Word64]
l
  where
        loop :: Int#
-> State# d
-> MutableByteArray# d
-> [Word64]
-> (# State# d, Array64 #)
loop Int#
_ State# d
st MutableByteArray# d
mb [] = forall {d}.
MutableByteArray# d -> State# d -> (# State# d, Array64 #)
freezeArray MutableByteArray# d
mb State# d
st
        loop Int#
i State# d
st MutableByteArray# d
mb ((W64# Word#
x):[Word64]
xs)
            | Int# -> Bool
booleanPrim (Int#
i Int# -> Int# -> Int#
==# Int#
n) = forall {d}.
MutableByteArray# d -> State# d -> (# State# d, Array64 #)
freezeArray MutableByteArray# d
mb State# d
st
            | Bool
otherwise =
                let !st' :: State# d
st' = forall d.
MutableByteArray# d -> Int# -> Word# -> State# d -> State# d
writeWord64Array# MutableByteArray# d
mb Int#
i Word#
x State# d
st
                 in Int#
-> State# d
-> MutableByteArray# d
-> [Word64]
-> (# State# d, Array64 #)
loop (Int#
i Int# -> Int# -> Int#
+# Int#
1#) State# d
st' MutableByteArray# d
mb [Word64]
xs
        freezeArray :: MutableByteArray# d -> State# d -> (# State# d, Array64 #)
freezeArray MutableByteArray# d
mb State# d
st =
            case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# d
mb State# d
st of
                (# State# d
st', ByteArray#
b #) -> (# State# d
st', ByteArray# -> Array64
Array64 ByteArray#
b #)
{-# NOINLINE array64 #-}

-- | Create a Mutable Array of Word32 of specific size from a list of Word32
mutableArray32 :: Int -> [Word32] -> IO MutableArray32
mutableArray32 :: Int -> [Word32] -> IO MutableArray32
mutableArray32 (I# Int#
n) [Word32]
l = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
    case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (Int#
n Int# -> Int# -> Int#
*# Int#
4#) Int#
4# State# RealWorld
s of
        (# State# RealWorld
s', MutableByteArray# RealWorld
mbarr #) -> Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> [Word32]
-> (# State# RealWorld, MutableArray32 #)
loop Int#
0# State# RealWorld
s' MutableByteArray# RealWorld
mbarr [Word32]
l
  where
        loop :: Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> [Word32]
-> (# State# RealWorld, MutableArray32 #)
loop Int#
_ State# RealWorld
st MutableByteArray# RealWorld
mb [] = (# State# RealWorld
st, MutableByteArray# RealWorld -> MutableArray32
MutableArray32 MutableByteArray# RealWorld
mb #)
        loop Int#
i State# RealWorld
st MutableByteArray# RealWorld
mb ((W32# Word32#
x):[Word32]
xs)
            | Int# -> Bool
booleanPrim (Int#
i Int# -> Int# -> Int#
==# Int#
n) = (# State# RealWorld
st, MutableByteArray# RealWorld -> MutableArray32
MutableArray32 MutableByteArray# RealWorld
mb #)
            | Bool
otherwise =
                let !st' :: State# RealWorld
st' = forall d.
MutableByteArray# d -> Int# -> Word32# -> State# d -> State# d
writeWord32Array# MutableByteArray# RealWorld
mb Int#
i Word32#
x State# RealWorld
st
                 in Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> [Word32]
-> (# State# RealWorld, MutableArray32 #)
loop (Int#
i Int# -> Int# -> Int#
+# Int#
1#) State# RealWorld
st' MutableByteArray# RealWorld
mb [Word32]
xs

-- | Create a Mutable Array of BE Word32 aliasing an Addr
mutableArray32FromAddrBE :: Int -> Addr# -> IO MutableArray32
mutableArray32FromAddrBE :: Int -> Addr# -> IO MutableArray32
mutableArray32FromAddrBE (I# Int#
n) Addr#
a = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s ->
    case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (Int#
n Int# -> Int# -> Int#
*# Int#
4#) Int#
4# State# RealWorld
s of
        (# State# RealWorld
s', MutableByteArray# RealWorld
mbarr #) -> Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> (# State# RealWorld, MutableArray32 #)
loop Int#
0# State# RealWorld
s' MutableByteArray# RealWorld
mbarr
  where
        loop :: Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> (# State# RealWorld, MutableArray32 #)
loop Int#
i State# RealWorld
st MutableByteArray# RealWorld
mb
            | Int# -> Bool
booleanPrim (Int#
i Int# -> Int# -> Int#
==# Int#
n) = (# State# RealWorld
st, MutableByteArray# RealWorld -> MutableArray32
MutableArray32 MutableByteArray# RealWorld
mb #)
            | Bool
otherwise             =
                let !st' :: State# RealWorld
st' = forall d.
MutableByteArray# d -> Int# -> Word32# -> State# d -> State# d
writeWord32Array# MutableByteArray# RealWorld
mb Int#
i (Word32# -> Word32#
be32Prim (Addr# -> Int# -> Word32#
indexWord32OffAddr# Addr#
a Int#
i)) State# RealWorld
st
                 in Int#
-> State# RealWorld
-> MutableByteArray# RealWorld
-> (# State# RealWorld, MutableArray32 #)
loop (Int#
i Int# -> Int# -> Int#
+# Int#
1#) State# RealWorld
st' MutableByteArray# RealWorld
mb

-- | freeze a Mutable Array of Word32 into a immutable Array of Word32
mutableArray32Freeze :: MutableArray32 -> IO Array32
mutableArray32Freeze :: MutableArray32 -> IO Array32
mutableArray32Freeze (MutableArray32 MutableByteArray# RealWorld
mb) = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
st ->
    case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# RealWorld
mb State# RealWorld
st of
        (# State# RealWorld
st', ByteArray#
b #) -> (# State# RealWorld
st', ByteArray# -> Array32
Array32 ByteArray#
b #)

-- | Read a Word8 from an Array
arrayRead8 :: Array8 -> Int -> Word8
arrayRead8 :: Array8 -> Int -> Word8
arrayRead8 (Array8 Addr#
a) (I# Int#
o) = Word8# -> Word8
W8# (Addr# -> Int# -> Word8#
indexWord8OffAddr# Addr#
a Int#
o)
{-# INLINE arrayRead8 #-}

-- | Read a Word32 from an Array
arrayRead32 :: Array32 -> Int -> Word32
arrayRead32 :: Array32 -> Int -> Word32
arrayRead32 (Array32 ByteArray#
b) (I# Int#
o) = Word32# -> Word32
W32# (ByteArray# -> Int# -> Word32#
indexWord32Array# ByteArray#
b Int#
o)
{-# INLINE arrayRead32 #-}

-- | Read a Word64 from an Array
arrayRead64 :: Array64 -> Int -> Word64
arrayRead64 :: Array64 -> Int -> Word64
arrayRead64 (Array64 ByteArray#
b) (I# Int#
o) = Word# -> Word64
W64# (ByteArray# -> Int# -> Word#
indexWord64Array# ByteArray#
b Int#
o)
{-# INLINE arrayRead64 #-}

-- | Read a Word32 from a Mutable Array of Word32
mutableArrayRead32 :: MutableArray32 -> Int -> IO Word32
mutableArrayRead32 :: MutableArray32 -> Int -> IO Word32
mutableArrayRead32 (MutableArray32 MutableByteArray# RealWorld
m) (I# Int#
o) = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Word32# #)
readWord32Array# MutableByteArray# RealWorld
m Int#
o State# RealWorld
s of (# State# RealWorld
s', Word32#
e #) -> (# State# RealWorld
s', Word32# -> Word32
W32# Word32#
e #)
{-# INLINE mutableArrayRead32 #-}

-- | Write a Word32 from a Mutable Array of Word32
mutableArrayWrite32 :: MutableArray32 -> Int -> Word32 -> IO ()
mutableArrayWrite32 :: MutableArray32 -> Int -> Word32 -> IO ()
mutableArrayWrite32 (MutableArray32 MutableByteArray# RealWorld
m) (I# Int#
o) (W32# Word32#
w) = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> let !s' :: State# RealWorld
s' = forall d.
MutableByteArray# d -> Int# -> Word32# -> State# d -> State# d
writeWord32Array# MutableByteArray# RealWorld
m Int#
o Word32#
w State# RealWorld
s in (# State# RealWorld
s', () #)
{-# INLINE mutableArrayWrite32 #-}

-- | Write into the Mutable Array of Word32 by combining through xor the current value and the new value.
--
-- > x[i] = x[i] xor value
mutableArrayWriteXor32 :: MutableArray32 -> Int -> Word32 -> IO ()
mutableArrayWriteXor32 :: MutableArray32 -> Int -> Word32 -> IO ()
mutableArrayWriteXor32 MutableArray32
m Int
o Word32
w =
    MutableArray32 -> Int -> IO Word32
mutableArrayRead32 MutableArray32
m Int
o forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word32
wOld -> MutableArray32 -> Int -> Word32 -> IO ()
mutableArrayWrite32 MutableArray32
m Int
o (Word32
wOld forall a. Bits a => a -> a -> a
`xor` Word32
w)
{-# INLINE mutableArrayWriteXor32 #-}