{-# LANGUAGE CPP #-}

-- |
-- Copyright:   (c) 2022 Andrew Lelechenko
--              (c) 2023 Pierre Le Marre
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Low-level routines for 'A.MArray' manipulations.
module Data.Text.Builder.Linear.Array (
  unsafeThaw,
  sizeofByteArray,
  isPinned,
  unsafeTile,
  unsafeReplicate,
) where

import Data.Text.Array qualified as A
import GHC.Exts (Int (..), isByteArrayPinned#, isTrue#, setByteArray#, sizeofByteArray#)
import GHC.ST (ST (..))

#if __GLASGOW_HASKELL__ >= 909
import GHC.Exts (unsafeThawByteArray#)
#else
import GHC.Exts (unsafeCoerce#)
#endif

unsafeThaw ∷ A.Array → ST s (A.MArray s)
#if __GLASGOW_HASKELL__ >= 909
unsafeThaw (A.ByteArray a) = ST $ \s# → case unsafeThawByteArray# a s# of
  (# s'#, ma #) -> (# s'#, A.MutableByteArray ma #)
#else
unsafeThaw :: forall s. Array -> ST s (MArray s)
unsafeThaw (A.ByteArray ByteArray#
a) = STRep s (MArray s) -> ST s (MArray s)
forall s a. STRep s a -> ST s a
ST (STRep s (MArray s) -> ST s (MArray s))
-> STRep s (MArray s) -> ST s (MArray s)
forall a b. (a -> b) -> a -> b
$ \State# s
s# →
  (# State# s
s#, MutableByteArray# s -> MArray s
forall s. MutableByteArray# s -> MArray s
A.MutableByteArray (ByteArray# -> MutableByteArray# s
forall a b. a -> b
unsafeCoerce# ByteArray#
a) #)
#endif

sizeofByteArray ∷ A.Array → Int
sizeofByteArray :: Array -> Int
sizeofByteArray (A.ByteArray ByteArray#
a) = Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
a)

isPinned ∷ A.Array → Bool
isPinned :: Array -> Bool
isPinned (A.ByteArray ByteArray#
a) = Int# -> Bool
isTrue# (ByteArray# -> Int#
isByteArrayPinned# ByteArray#
a)

-- | Replicate an ASCII character
--
-- __Warning:__ it is the responsibility of the caller to ensure that the 'Int'
-- is a valid ASCII character.
unsafeReplicate
  ∷ A.MArray s
  -- ^ Mutable array
  → Int
  -- ^ Offset
  → Int
  -- ^ Count
  → Int
  -- ^ ASCII character
  → ST s ()
unsafeReplicate :: forall s. MArray s -> Int -> Int -> Int -> ST s ()
unsafeReplicate (A.MutableByteArray MutableByteArray# s
dst#) (I# Int#
dstOff#) (I# Int#
count#) (I# Int#
w#) =
  STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (\State# s
s# → (# MutableByteArray# s -> Int# -> Int# -> Int# -> State# s -> State# s
forall d.
MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
setByteArray# MutableByteArray# s
dst# Int#
dstOff# Int#
count# Int#
w# State# s
s#, () #))
{-# INLINE unsafeReplicate #-}

-- | Duplicate a portion of an array in-place.
--
-- Example of use:
--
-- @
-- -- Write @count@ times the char @c@
-- let cLen = utf8Length c; totalLen = cLen * count
-- in unsafeWrite dst dstOff ch *> 'unsafeTile' dst dstOff totalLen cLen
-- @
unsafeTile
  ∷ A.MArray s
  -- ^ Mutable array
  → Int
  -- ^ Start of the portion to duplicate
  → Int
  -- ^ Total length of the duplicate
  → Int
  -- ^ Length of the portion to duplicate
  → ST s ()
unsafeTile :: forall s. MArray s -> Int -> Int -> Int -> ST s ()
unsafeTile MArray s
dest Int
destOff Int
totalLen = Int -> ST s ()
go
  where
    -- Adapted from Data.Text.Array.tile
    go :: Int -> ST s ()
go Int
l
      | Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
totalLen = MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
A.copyM MArray s
dest (Int
destOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l) MArray s
dest Int
destOff (Int
totalLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
      | Bool
otherwise = MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
A.copyM MArray s
dest (Int
destOff Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l) MArray s
dest Int
destOff Int
l ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
go (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
l)
{-# INLINE unsafeTile #-}