{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module: Data.ByteString.Base16
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Pure base16 encoding and decoding of strict bytestrings.

module Data.ByteString.Base16 (
    encode
  , decode
  ) where

import qualified Data.Bits as B
import Data.Bits ((.&.), (.|.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base16.Arm as Arm
import qualified Data.ByteString.Internal as BI
import Data.Word (Word8, Word16)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import Foreign.Storable (peekElemOff, pokeElemOff)
import System.IO.Unsafe (unsafeDupablePerformIO)

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

-- 512-byte table.  Bytes [2k] and [2k+1] are the two lowercase ASCII
-- hex characters representing the value k.  All-ASCII content means
-- the bytestring 'IsString' rule rewrites this to 'unsafePackAddress'
-- and the bytes live in static rodata.
enc_tab :: BS.ByteString
enc_tab :: ByteString
enc_tab =
  ByteString
"000102030405060708090a0b0c0d0e0f\
  \101112131415161718191a1b1c1d1e1f\
  \202122232425262728292a2b2c2d2e2f\
  \303132333435363738393a3b3c3d3e3f\
  \404142434445464748494a4b4c4d4e4f\
  \505152535455565758595a5b5c5d5e5f\
  \606162636465666768696a6b6c6d6e6f\
  \707172737475767778797a7b7c7d7e7f\
  \808182838485868788898a8b8c8d8e8f\
  \909192939495969798999a9b9c9d9e9f\
  \a0a1a2a3a4a5a6a7a8a9aaabacadaeaf\
  \b0b1b2b3b4b5b6b7b8b9babbbcbdbebf\
  \c0c1c2c3c4c5c6c7c8c9cacbcccdcecf\
  \d0d1d2d3d4d5d6d7d8d9dadbdcdddedf\
  \e0e1e2e3e4e5e6e7e8e9eaebecedeeef\
  \f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"
{-# NOINLINE enc_tab #-}

-- 256-byte table.  Index by an ASCII byte to obtain its nibble; valid
-- hex chars ('0'..'9', 'a'..'f', 'A'..'F') map to 0x10..0x1f, every
-- other byte maps to 0x20.
--
-- The encoding is chosen so the literal is strictly ASCII and contains
-- no embedded NUL, which is what the bytestring 'IsString' rule needs
-- to rewrite it into 'unsafePackAddress' (cf. 'enc_tab') — the bytes
-- end up in static rodata, with no CAF allocation.
--
-- The 0x20 sentinel is distinguished by bit 5; no value 0x10..0x1f
-- carries that bit, so 'decode' OR-folds every lookup into an
-- accumulator and tests 'acc .&. 0x20 == 0' once at the end.  The
-- output byte is '(n0 `shiftL` 4) .|. (n1 .&. 0x0f)': in 'Word8' the
-- shift naturally drops bit 4, and the mask isolates the low nibble.
dec_tab :: BS.ByteString
dec_tab :: ByteString
dec_tab =
  ByteString
"\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x20\x20\x20\x20\x20\x20\
  \\x20\x1a\x1b\x1c\x1d\x1e\x1f\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x1a\x1b\x1c\x1d\x1e\x1f\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\
  \\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20\x20"
{-# NOINLINE dec_tab #-}

-- | Encode a base256 'ByteString' as base16.
--
--   Uses ARM NEON extensions when available, otherwise a pure
--   Haskell scalar loop.
--
--   >>> encode "hello world"
--   "68656c6c6f20776f726c64"
encode :: BS.ByteString -> BS.ByteString
encode :: ByteString -> ByteString
encode ByteString
bs
  | Bool
Arm.base16_arm_available = ByteString -> ByteString
Arm.encode ByteString
bs
  | Bool
otherwise = ByteString -> ByteString
encode_scalar ByteString
bs
{-# INLINABLE encode #-}

-- | Decode a base16 'ByteString' to base256.
--
--   Uses ARM NEON extensions when available, otherwise a pure
--   Haskell scalar loop.  Invalid inputs (including odd-length
--   inputs) will produce 'Nothing'.
--
--   >>> decode "68656c6c6f20776f726c64"
--   Just "hello world"
--   >>> decode "068656c6c6f20776f726c64" -- odd-length
--   Nothing
decode :: BS.ByteString -> Maybe BS.ByteString
decode :: ByteString -> Maybe ByteString
decode ByteString
bs
  | Bool
Arm.base16_arm_available = ByteString -> Maybe ByteString
Arm.decode ByteString
bs
  | Bool
otherwise = ByteString -> Maybe ByteString
decode_scalar ByteString
bs
{-# INLINABLE decode #-}

encode_scalar :: BS.ByteString -> BS.ByteString
encode_scalar :: ByteString -> ByteString
encode_scalar (BI.PS ForeignPtr Word8
sfp Int
soff Int
l) =
  case ByteString
enc_tab of
    BI.PS ForeignPtr Word8
tfp Int
toff Int
_ ->
      Int -> (Ptr Word8 -> IO ()) -> ByteString
BI.unsafeCreate (Int
l Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
1) ((Ptr Word8 -> IO ()) -> ByteString)
-> (Ptr Word8 -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dst ->
        ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp0 ->
        ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
tfp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
tp0 -> do
          -- read 'enc_tab' and write 'dst' as 'Word16' pairs.  The
          -- two-byte block at 'enc_tab[2*b]' and the two-byte block
          -- at 'dst[2*i]' share the same byte layout in memory, so
          -- this is endianness-safe: we never inspect the numerical
          -- value of the 'Word16', we just shuffle 16 bits between
          -- two locations.
          let !sp :: Ptr Word8
sp = Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff :: Ptr Word8
              !tp :: Ptr Word16
tp = Ptr Word8
tp0 Ptr Word8 -> Int -> Ptr Word16
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
toff :: Ptr Word16
              !dp :: Ptr Word16
dp = Ptr Word8 -> Ptr Word16
forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
dst        :: Ptr Word16
              loop :: Int -> IO ()
loop !Int
i
                | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l    = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    b <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp Int
i
                    w <- peekElemOff tp (fi b)
                    pokeElemOff dp i (w :: Word16)
                    loop (i + 1)
          Int -> IO ()
loop Int
0

decode_scalar :: BS.ByteString -> Maybe BS.ByteString
decode_scalar :: ByteString -> Maybe ByteString
decode_scalar (BI.PS ForeignPtr Word8
sfp Int
soff Int
l)
  | Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
B.testBit Int
l Int
0 = Maybe ByteString
forall a. Maybe a
Nothing
  | Bool
otherwise = case ByteString
dec_tab of
      BI.PS ForeignPtr Word8
tfp Int
toff Int
_ -> IO (Maybe ByteString) -> Maybe ByteString
forall a. IO a -> a
unsafeDupablePerformIO (IO (Maybe ByteString) -> Maybe ByteString)
-> IO (Maybe ByteString) -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ do
        let !n :: Int
n = Int
l Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftR` Int
1
        fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
BI.mallocByteString Int
n
        ok <- withForeignPtr fp  $ \Ptr Word8
dst ->
              ForeignPtr Word8 -> (Ptr Word8 -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO Bool) -> IO Bool)
-> (Ptr Word8 -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp0 ->
              ForeignPtr Word8 -> (Ptr Word8 -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
tfp ((Ptr Word8 -> IO Bool) -> IO Bool)
-> (Ptr Word8 -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
tp0 -> do
                let !sp :: Ptr Word8
sp = Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff :: Ptr Word8
                    !tp :: Ptr Word8
tp = Ptr Word8
tp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
toff :: Ptr Word8
                    loop :: Int -> Word8 -> IO Bool
loop !Int
i !Word8
acc
                      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n    =
                          Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$! Word8
acc Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x20 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0
                      | Bool
otherwise = do
                          let !o :: Int
o = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
1
                          c0 <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
sp  Int
o
                          c1 <- peekElemOff sp (o + 1)
                          n0 <- peekElemOff tp (fi c0)
                          n1 <- peekElemOff tp (fi c1)
                          let !b = (Word8
n0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`B.shiftL` Int
4)
                               Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. (Word8
n1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x0f)
                          pokeElemOff dst i b
                          loop (i + 1) (acc .|. n0 .|. n1)
                Int -> Word8 -> IO Bool
loop Int
0 Word8
0
        pure $! if ok then Just (BI.PS fp 0 n) else Nothing