Arm.hs (2852B)
1 {-# OPTIONS_HADDOCK hide #-} 2 {-# LANGUAGE BangPatterns #-} 3 4 -- | 5 -- Module: Data.ByteString.Base64.Arm 6 -- Copyright: (c) 2026 Jared Tobin 7 -- License: MIT 8 -- Maintainer: Jared Tobin <jared@ppad.tech> 9 -- 10 -- ARM NEON support for base64 encoding and decoding. 11 12 module Data.ByteString.Base64.Arm ( 13 base64_arm_available 14 , encode 15 , decode 16 ) where 17 18 import qualified Data.Bits as B 19 import Data.Bits ((.&.)) 20 import qualified Data.ByteString as BS 21 import qualified Data.ByteString.Internal as BI 22 import Data.Word (Word8) 23 import Foreign.C.Types (CInt(..), CSize(..)) 24 import Foreign.ForeignPtr (withForeignPtr) 25 import Foreign.Ptr (Ptr, plusPtr) 26 import Foreign.Storable (peekElemOff) 27 import System.IO.Unsafe (unsafeDupablePerformIO) 28 29 -- ffi ------------------------------------------------------------------------ 30 31 foreign import ccall unsafe "base64_encode_arm" 32 c_base64_encode :: Ptr Word8 -> Ptr Word8 -> CSize -> IO () 33 34 foreign import ccall unsafe "base64_decode_arm" 35 c_base64_decode :: Ptr Word8 -> Ptr Word8 -> CSize -> CSize -> IO CInt 36 37 foreign import ccall unsafe "base64_arm_available" 38 c_base64_arm_available :: IO CInt 39 40 -- utilities ------------------------------------------------------------------ 41 42 fi :: (Integral a, Num b) => a -> b 43 fi = fromIntegral 44 {-# INLINE fi #-} 45 46 -- api ------------------------------------------------------------------------ 47 48 -- | Are ARM NEON extensions available? 49 base64_arm_available :: Bool 50 base64_arm_available = 51 unsafeDupablePerformIO c_base64_arm_available /= 0 52 {-# NOINLINE base64_arm_available #-} 53 54 -- | Encode a base256 'ByteString' as base64 using NEON. 55 encode :: BS.ByteString -> BS.ByteString 56 encode (BI.PS sfp soff l) = 57 BI.unsafeCreate ((l + 2) `quot` 3 * 4) $ \dst -> 58 withForeignPtr sfp $ \sp0 -> 59 c_base64_encode (sp0 `plusPtr` soff) dst (fi l) 60 61 -- | Decode a base64 'ByteString' to base256 using NEON. Returns 62 -- 'Nothing' on malformed input. 63 decode :: BS.ByteString -> Maybe BS.ByteString 64 decode (BI.PS sfp soff l) 65 | l == 0 = Just BS.empty 66 | l .&. 0x03 /= 0 = Nothing 67 | otherwise = unsafeDupablePerformIO $ 68 withForeignPtr sfp $ \sp0 -> do 69 let !sp = sp0 `plusPtr` soff :: Ptr Word8 70 c_pre <- peekElemOff sp (l - 2) 71 c_end <- peekElemOff sp (l - 1) 72 let !pad_pre = c_pre == 0x3D 73 !pad_end = c_end == 0x3D 74 if pad_pre && not pad_end 75 then pure Nothing 76 else do 77 let !pad = (if pad_pre then 2 else if pad_end then 1 else 0) 78 :: Int 79 !nfull = l `B.shiftR` 2 80 !outlen = nfull * 3 - pad 81 fp <- BI.mallocByteString outlen 82 ok <- withForeignPtr fp $ \dst -> 83 c_base64_decode sp dst (fi l) (fi outlen) 84 pure $! if ok /= 0 85 then Just (BI.PS fp 0 outlen) 86 else Nothing