{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}

module Country.Unexposed.Util
  ( newZeroedByteArray
  , mapTextArray
  , charToWord8
  , word16ToChar
  , word16ToInt
  , timesTwo
  , timesThree
  , half
  ) where

import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Bits (unsafeShiftL, unsafeShiftR)
import Data.Char (chr, ord)
import Data.Primitive.ByteArray (MutableByteArray, fillByteArray, newByteArray)
import Data.Word (Word16, Word8)
import GHC.Exts (sizeofByteArray#)
import GHC.Int (Int (I#))

import qualified Data.Text.Array as TA

newZeroedByteArray :: (PrimMonad m) => Int -> m (MutableByteArray (PrimState m))
newZeroedByteArray :: forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newZeroedByteArray Int
len = do
  MutableByteArray (PrimState m)
arr <- Int -> m (MutableByteArray (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
len
  let arrStart :: Int
arrStart = Int
0
  let zeroByte :: Word8
zeroByte = Word8
0
  MutableByteArray (PrimState m) -> Int -> Int -> Word8 -> m ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Int -> Word8 -> m ()
fillByteArray MutableByteArray (PrimState m)
arr Int
arrStart Int
len Word8
zeroByte
  MutableByteArray (PrimState m)
-> m (MutableByteArray (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableByteArray (PrimState m)
arr
{-# INLINE newZeroedByteArray #-}

mapTextArray :: (Char -> Char) -> TA.Array -> TA.Array
mapTextArray :: (Char -> Char) -> Array -> Array
mapTextArray Char -> Char
f a :: Array
a@(TA.ByteArray ByteArray#
inner) = (forall s. ST s (MArray s)) -> Array
TA.run ((forall s. ST s (MArray s)) -> Array)
-> (forall s. ST s (MArray s)) -> Array
forall a b. (a -> b) -> a -> b
$ do
  let len :: Int
len = Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
inner)
  MArray s
m <- Int -> ST s (MArray s)
forall s. Int -> ST s (MArray s)
TA.new Int
len
  Int -> MArray s -> Int -> Array -> Int -> ST s ()
forall s. Int -> MArray s -> Int -> Array -> Int -> ST s ()
TA.copyI Int
len MArray s
m Int
0 Array
a Int
0
  let go :: Int -> ST s ()
go !Int
ix =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
          then do
            MArray s -> Int -> Word8 -> ST s ()
forall s. MArray s -> Int -> Word8 -> ST s ()
TA.unsafeWrite MArray s
m Int
ix (Char -> Word8
charToWord8 (Char -> Char
f (Word8 -> Char
word8ToChar (Array -> Int -> Word8
TA.unsafeIndex Array
a Int
ix))))
            Int -> ST s ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          else () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  Int -> ST s ()
go Int
0
  MArray s -> ST s (MArray s)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s
m
{-# INLINE mapTextArray #-}

charToWord8 :: Char -> Word8
charToWord8 :: Char -> Word8
charToWord8 = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Char -> Int) -> Char -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord
{-# INLINE charToWord8 #-}

word8ToChar :: Word8 -> Char
word8ToChar :: Word8 -> Char
word8ToChar = Int -> Char
chr (Int -> Char) -> (Word8 -> Int) -> Word8 -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE word8ToChar #-}

word16ToChar :: Word16 -> Char
word16ToChar :: Word16 -> Char
word16ToChar = Int -> Char
chr (Int -> Char) -> (Word16 -> Int) -> Word16 -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE word16ToChar #-}

word16ToInt :: Word16 -> Int
word16ToInt :: Word16 -> Int
word16ToInt = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE word16ToInt #-}

timesTwo :: Int -> Int
timesTwo :: Int -> Int
timesTwo Int
x = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
x Int
1
{-# INLINE timesTwo #-}

timesThree :: Int -> Int
timesThree :: Int -> Int
timesThree Int
x = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
timesTwo Int
x
{-# INLINE timesThree #-}

half :: Int -> Int
half :: Int -> Int
half Int
x = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
x Int
1
{-# INLINE half #-}