Utils.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
-- |
-- Module       : Data.ByteString.Base32.Internal.Utils
-- Copyright    : (c) 2020-2026 Emily Pillmore
-- License      : BSD-style
--
-- Maintainer   : Emily Pillmore <emilypi@cohomolo.gy>
-- Stability    : Experimental
-- Portability  : portable
--
-- Shared allocation, byte-order, table lookup, rechunking, and padding utilities.
--
-- @since 1.0
--
module Data.ByteString.Base32.Internal.Utils
( aix
, peekWord32BE
, peekWord64BE
, reChunkN
, w32
, w64
, w64_32
, encodedLength
, writeNPlainPtrBytes
, validateLastNPads
) where


import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Text (Text)
import qualified Data.Text as Text

import Foreign.Ptr
import Foreign.Storable

import GHC.ByteOrder
import GHC.Exts
import GHC.Word

import System.IO.Unsafe
import Foreign.Marshal.Alloc (mallocBytes)


-- | Read 'Word8' index off alphabet addr
--
-- @since 1.0
--
aix :: Word8 -> Addr# -> Word8
aix w8 alpha = W8# (indexWord8OffAddr# alpha i)
  where
    !(I# i) = fromIntegral w8
{-# INLINE aix #-}

-- | Widen a 'Word8' to a 'Word32'.
--
-- @since 1.0
--
w32 :: Word8 -> Word32
w32 = fromIntegral
{-# INLINE w32 #-}

-- | Widen a 'Word32' to a 'Word64'.
--
-- @since 1.0
--
w64_32 :: Word32 -> Word64
w64_32 = fromIntegral
{-# INLINE w64_32 #-}

-- | Widen a 'Word8' to a 'Word64'.
--
-- @since 1.0
--
w64 :: Word8 -> Word64
w64 = fromIntegral
{-# INLINE w64 #-}

-- | Allocate and fill @n@ bytes with some data
--
-- @since 1.0
--
writeNPlainPtrBytes
  :: Storable a
  => Int
  -> [a]
  -> Ptr a
writeNPlainPtrBytes !n as = unsafeDupablePerformIO $ do
    p <- mallocBytes n
    go p as
    return p
  where
    go !_ [] = return ()
    go !p (x:xs) = poke p x >> go (plusPtr p 1) xs
{-# INLINE writeNPlainPtrBytes #-}

-- | Read a 32-bit word in big-endian byte order.
--
-- @since 1.0
--
peekWord32BE :: Ptr Word32 -> IO Word32
peekWord32BE p = case targetByteOrder of
  LittleEndian -> byteSwap32 <$> peek p
  BigEndian -> peek p
{-# inline peekWord32BE #-}

-- | Read a 64-bit word in big-endian byte order.
--
-- @since 1.0
--
peekWord64BE :: Ptr Word64 -> IO Word64
peekWord64BE p = case targetByteOrder of
  LittleEndian -> byteSwap64 <$> peek p
  BigEndian -> peek p
{-# inline peekWord64BE #-}

-- | Rechunk a list of bytestrings in multiples of @n@
--
-- @since 1.0
--
reChunkN :: Int -> [ByteString] -> [ByteString]
reChunkN n = go
  where
    go [] = []
    go (b:bs) = case divMod (BS.length b) n of
      (_, 0) -> b : go bs
      (d, _) -> case BS.splitAt (d * n) b of
        ~(h, t) -> h : accum t bs

    accum acc [] = [acc]
    accum acc (c:cs) =
      case BS.splitAt (n - BS.length acc) c of
        ~(h, t) ->
          let acc' = BS.append acc h
          in if BS.length acc' == n
             then
               let cs' = if BS.null t then cs else t : cs
               in acc' : go cs'
             else accum acc' cs
{-# INLINE reChunkN #-}

-- | Padded output capacity, using checked integer arithmetic.
--
-- @since 1.0
--
encodedLength :: Int -> Int
encodedLength n
  | n < 0 || blocks > maxBound `quot` 8 = error "Base32 output is too large"
  | otherwise = blocks * 8
  where
    blocks = n `quot` 5 + if n `rem` 5 == 0 then 0 else 1
{-# INLINE encodedLength #-}

-- | Reject '=' in the final @n@ bytes before running the decoder.
-- The caller selects @n@ from the input length and padding policy, preventing
-- existing padding from being accepted when missing padding is added.
--
-- @since 1.0
--
validateLastNPads
    :: Int
    -> ByteString
    -> IO (Either Text ByteString)
    -> Either Text ByteString
validateLastNPads n bs io
    | BS.elem 0x3d (BS.drop (BS.length bs - n) bs) =
        Left (Text.pack "Base32-encoded bytestring has invalid padding")
    | otherwise = unsafeDupablePerformIO io
{-# INLINE validateLastNPads #-}