Skip to content

Commit

Permalink
Merge pull request #4 from commercialhaskell/fix1
Browse files Browse the repository at this point in the history
Fix #1 Enforce that the Word64 values used are little-endian.
  • Loading branch information
mpilgrem authored Oct 21, 2024
2 parents ec6f717 + 08d19aa commit e213b2d
Showing 1 changed file with 20 additions and 4 deletions.
24 changes: 20 additions & 4 deletions src/Data/StaticBytes.hs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ import qualified Data.Vector.Unboxed.Base as VU
import Foreign.ForeignPtr ( ForeignPtr, withForeignPtr )
import Foreign.Ptr ( Ptr, castPtr )
import Foreign.Storable ( Storable (..) )
import GHC.ByteOrder ( ByteOrder (..), targetByteOrder )
import RIO hiding ( words )
import System.IO.Unsafe ( unsafePerformIO )

Expand Down Expand Up @@ -79,27 +80,31 @@ class DynamicBytes dbytes where
lengthD :: dbytes -> Int
-- Yeah, it looks terrible to use a list here, but fusion should kick in
withPeekD :: dbytes -> ((Int -> IO Word64) -> IO a) -> IO a
-- ^ This assumes that the Word64 values are all little-endian.
-- | May throw a runtime exception if invariants are violated!
fromWordsD :: Int -> [Word64] -> dbytes
-- ^ This assumes that the Word64 values are all little-endian.

fromWordsForeign ::
(ForeignPtr a -> Int -> b)
-> Int
-> [Word64]
-- ^ The Word64 values are assumed to be little-endian.
-> b
fromWordsForeign wrapper len words0 = unsafePerformIO $ do
fptr <- B.mallocByteString len
withForeignPtr fptr $ \ptr -> do
let loop _ [] = pure ()
loop off (w:ws) = do
pokeElemOff (castPtr ptr) off w
pokeElemOff (castPtr ptr) off (fromLE64 w)
loop (off + 1) ws
loop 0 words0
pure $ wrapper fptr len

withPeekForeign ::
(ForeignPtr a, Int, Int)
-> ((Int -> IO Word64) -> IO b)
-- ^ The Word64 values are assumed to be little-endian.
-> IO b
withPeekForeign (fptr, off, len) inner =
withForeignPtr fptr $ \ptr -> do
Expand All @@ -113,7 +118,7 @@ withPeekForeign (fptr, off, len) inner =
let w64' = shiftL (fromIntegral w8) (i * 8) .|. w64
loop w64' (i + 1)
loop 0 0
| otherwise = peekByteOff ptr (off + off')
| otherwise = toLE64 <$> peekByteOff ptr (off + off')
inner f

instance DynamicBytes B.ByteString where
Expand All @@ -133,7 +138,7 @@ instance word8 ~ Word8 => DynamicBytes (VP.Vector word8) where
let loop _ [] =
VP.Vector 0 len <$> BA.unsafeFreezeByteArray ba
loop i (w:ws) = do
BA.writeByteArray ba i w
BA.writeByteArray ba i (fromLE64 w)
loop (i + 1) ws
loop 0 words0
withPeekD (VP.Vector off len ba) inner = do
Expand All @@ -147,7 +152,8 @@ instance word8 ~ Word8 => DynamicBytes (VP.Vector word8) where
let w64' = shiftL (fromIntegral w8) (i * 8) .|. w64
loop w64' (i + 1)
loop 0 0
| otherwise = pure $ BA.indexByteArray ba (off + (off' `div` 8))
| otherwise = pure $
toLE64 $ BA.indexByteArray ba (off + (off' `div` 8))
inner f

instance word8 ~ Word8 => DynamicBytes (VU.Vector word8) where
Expand Down Expand Up @@ -248,3 +254,13 @@ fromStatic ::
=> sbytes
-> dbytes
fromStatic = fromWordsD (lengthS (Nothing :: Maybe sbytes)) . ($ []) . toWordsS

-- | Convert a 64 bit value in CPU endianess to little endian.
toLE64 :: Word64 -> Word64
toLE64 = case targetByteOrder of
BigEndian -> byteSwap64
LittleEndian -> id

-- | Convert a little endian 64 bit value to CPU endianess.
fromLE64 :: Word64 -> Word64
fromLE64 = toLE64

0 comments on commit e213b2d

Please sign in to comment.