{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE BangPatterns #-}

-- |
-- Module: Data.ByteString.Base16.Arm
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- ARM NEON support for base16 encoding and decoding.

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

import qualified Data.Bits as B
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI
import Data.Word (Word8)
import Foreign.C.Types (CInt(..), CSize(..))
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, plusPtr)
import System.IO.Unsafe (unsafeDupablePerformIO)

-- ffi ------------------------------------------------------------------------

foreign import ccall unsafe "base16_encode_arm"
  c_base16_encode :: Ptr Word8 -> Ptr Word8 -> CSize -> IO ()

foreign import ccall unsafe "base16_decode_arm"
  c_base16_decode :: Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt

foreign import ccall unsafe "base16_arm_available"
  c_base16_arm_available :: IO CInt

-- utilities ------------------------------------------------------------------

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

-- api ------------------------------------------------------------------------

-- | Are ARM NEON extensions available?
base16_arm_available :: Bool
base16_arm_available :: Bool
base16_arm_available =
  IO CInt -> CInt
forall a. IO a -> a
unsafeDupablePerformIO IO CInt
c_base16_arm_available CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
/= CInt
0
{-# NOINLINE base16_arm_available #-}

-- | Encode a base256 'ByteString' as base16 using NEON.
encode :: BS.ByteString -> BS.ByteString
encode :: ByteString -> ByteString
encode (BI.PS ForeignPtr Word8
sfp Int
soff Int
l) =
  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 ->
      Ptr Word8 -> Ptr Word8 -> CSize -> IO ()
c_base16_encode (Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff) Ptr Word8
dst (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fi Int
l)

-- | Decode a base16 'ByteString' to base256 using NEON.  Returns
--   'Nothing' on odd-length or otherwise invalid input.
decode :: BS.ByteString -> Maybe BS.ByteString
decode :: ByteString -> Maybe ByteString
decode (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 = 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 CInt) -> IO CInt
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO CInt) -> IO CInt)
-> (Ptr Word8 -> IO CInt) -> IO CInt
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp0 ->
              Ptr Word8 -> Ptr Word8 -> CSize -> IO CInt
c_base16_decode (Ptr Word8
sp0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
soff) Ptr Word8
dst (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fi Int
n)
      pure $! if ok /= 0 then Just (BI.PS fp 0 n) else Nothing