{-# 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) = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s# →
  (# State# s
s#, forall s. MutableByteArray# s -> MutableByteArray s
A.MutableByteArray (unsafeCoerce# :: 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#) =
  forall s a. STRep s a -> ST s a
ST (\State# s
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 forall a. Num a => a -> a -> a
* Int
l forall a. Ord a => a -> a -> Bool
> Int
totalLen = forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
A.copyM MArray s
dest (Int
destOff forall a. Num a => a -> a -> a
+ Int
l) MArray s
dest Int
destOff (Int
totalLen forall a. Num a => a -> a -> a
- Int
l)
      | Bool
otherwise = forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
A.copyM MArray s
dest (Int
destOff forall a. Num a => a -> a -> a
+ Int
l) MArray s
dest Int
destOff Int
l forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
go (Int
2 forall a. Num a => a -> a -> a
* Int
l)
{-# INLINE unsafeTile #-}