{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module Data.LLVM.BitCode.GetBits (
GetBits
, runGetBits
, fixed, align32bits
, bytestring
, label
, isolate
, try
, skip
) where
import Data.LLVM.BitCode.BitString
import Control.Applicative ( Alternative(..) )
import Control.Monad ( MonadPlus(..) )
import Data.Bits ( shiftR, shiftL, (.&.), (.|.) )
import Data.ByteString ( ByteString )
import qualified Data.ByteString as BS
import GHC.Exts
import GHC.Word
#if !MIN_VERSION_base(4,13,0)
import Control.Monad.Fail ( MonadFail )
import qualified Control.Monad.Fail
#endif
-- Bit-level Parsing -----------------------------------------------------------
newtype GetBits a =
GetBits { unGetBits :: BitPosition -> BS.ByteString
-> (# BitsGetter a, BitPosition #)
}
type BitPosition = (# Int#, Int# #) -- (# current bit pos, maximum bit pos #)
type BitsGetter a = Either String a -- Left is fail
-- | Run a @GetBits@ action, returning its value, and the number of bits offset
-- into the next byte of the stream.
runGetBits :: GetBits a -> ByteString -> Either String a
runGetBits m bs =
let !startPos# = (# 0#, bitCount# $ bytesToBits $ Bytes' $ BS.length bs #)
!(# g, _ #) = unGetBits m startPos# bs
in g
instance Functor GetBits where
{-# INLINE fmap #-}
fmap f m = GetBits $
\ !pos# inp -> let !(# b, n# #) = unGetBits m pos# inp
in (# f <$> b, n# #)
instance Applicative GetBits where
{-# INLINE pure #-}
pure x = GetBits $ \ !pos# _ -> (# pure x, pos# #)
{-# INLINE (<*>) #-}
f <*> x =
GetBits $ \ !pos# inp ->
let !(# g, n# #) = unGetBits f pos# inp
in case g of
Right g' ->
let !(# y, m# #) = unGetBits x n# inp
in case y of
Right y' -> (# Right $ g' y', m# #)
Left e -> (# Left e, m# #)
Left e -> (# Left e, n# #)
instance Monad GetBits where
{-# INLINE return #-}
return = pure
{-# INLINE (>>=) #-}
m >>= f = GetBits $ \ !pos# inp ->
let !(# g, n# #) = unGetBits m pos# inp
!(# gr, nr# #) = case g of
Left e -> (# Left e, n# #)
Right a -> unGetBits (f a) n# inp
in (# gr, nr# #)
#if !MIN_VERSION_base(4,13,0)
{-# INLINE fail #-}
fail e = GetBits $ \ p _ -> (# Left e, p #)
#endif
instance MonadFail GetBits where
{-# INLINE fail #-}
fail e = GetBits $ \ p _ -> (# Left e, p #)
instance Alternative GetBits where
{-# INLINE empty #-}
empty = GetBits $ \ p _ -> (# Left "GetBits is empty!", p #)
{-# INLINE (<|>) #-}
a <|> b = GetBits
$ \ !pos# inp ->
let !r@(# g, _ #) = unGetBits a pos# inp
in case g of
Right _ -> r
Left _ -> unGetBits b pos# inp
instance MonadPlus GetBits where
{-# INLINE mzero #-}
mzero = empty
{-# INLINE mplus #-}
mplus = (<|>)
-- | Extracts an Integer value of the specified number of bits from a ByteString,
-- starting at the indicated bit position (fails with Left if the range to
-- extract is not valid... i.e. > bitLimit). Returns the Integer value along
-- with the bit position following the extraction.
-- There are two implementations: one builds an integer value by shifting bits,
-- then shifts and masks the result to get the final value. The other uses
-- unlifted values and avoids the final shift by being smarter about individual
-- compositions. Their functionality should be identical, but it may be easier
-- to debug the first.
_extractFromByteString' :: NumBits {-^ the last bit accessible in the ByteString -}
-> NumBits {-^ the bit to start extraction at -}
-> NumBits {-^ the number of bits to extract -}
-> ByteString {-^ the ByteString to extract from -}
-> Either String (Int, NumBits)
_extractFromByteString' bitLimit startBit numBits bs =
let Bytes' s8 = fst (bitsToBytes startBit)
Bytes' r8 = fst (bitsToBytes numBits)
rcnt = r8 + 2 -- 2 == pre-shift overflow byte on either side
-- Extract the relevant bits from the ByteCode, with padding to byte
-- boundaries into ws.
ws = BS.take rcnt $ BS.drop s8 bs
-- Combine the extracted bytes into an Integer value in wi.
wi = BS.foldr (\w a -> a `shiftL` 8 .|. fromIntegral w) (0::Int) ws
-- Mask is 0-bit based set of bits wanted in the result
mask = ((1::Int) `shiftL` bitCount numBits) - 1
-- Shift the desired value down to byte alignment and then discard any
-- excess high bits.
vi = wi `shiftR` (bitCount startBit .&. 7) .&. mask
updPos = addBitCounts startBit numBits
in if updPos > bitLimit
then Left ("Attempt to read bits past limit (newPos="
<> show updPos <> ", limit=" <> show bitLimit <> ")"
)
else Right (vi, updPos)
extractFromByteString :: Int# {-^ the last bit accessible in the ByteString -}
-> Int# {-^ the bit to start extraction at -}
-> Int# {-^ the number of bits to extract -}
-> ByteString {-^ the ByteString to extract from -}
-> Either String (() -> (# Int#, Int# #))
extractFromByteString !bitLim# !sBit# !nbits# bs =
if isTrue# ((1# `uncheckedIShiftL#` (nbits#)) /=# 0#)
-- (nbits# -# 1#) above would allow 64-bit value extraction, but this
-- function cannot actually support a size of 64, because Int# is signed,
-- so it doesn't properly use the high bit in numeric operations. This
-- seems to be OK at this point because LLVM bitcode does not attempt to
-- encode actual 64-bit values.
then
let !updPos# = sBit# +# nbits#
in if isTrue# (updPos# <=# bitLim#)
then
let !s8# = sBit# `uncheckedIShiftRL#` 3#
!hop# = sBit# `andI#` 7#
!r8# = ((hop# +# nbits# +# 7#) `uncheckedIShiftRL#` 3#)
!mask# = (1# `uncheckedIShiftL#` nbits#) -# 1#
-- Here, s8# is the size in 8-bit bytes, hop# is the number of
-- bits shifted from the byte boundary, r8# is the rounded number
-- of bytes actually needed to retrieve to get the value to
-- account for shifting, and mask# is the mask for the final
-- target set of bits after shifting.
#if MIN_VERSION_base(4,16,0)
word8ToInt !w8# = word2Int# (word8ToWord# w8#)
#else
-- technically #if !MIN_VERSION_ghc_prim(0,8,0), for GHC 9.2, but
-- since ghc_prim isn't a direct dependency and is re-exported
-- from base, this define needs to reference the base version.
word8ToInt = word2Int#
#endif
-- getB# gets a value from a byte starting at bit0 of the byte
getB# :: Int# -> Int#
getB# !i# =
case i# of
0# -> let !(W8# w#) = bs `BS.index` (I# s8#)
in word8ToInt w#
_ -> let !(W8# w#) = (bs `BS.index` (I# (s8# +# i#)))
in (word8ToInt w#) `uncheckedIShiftL#` (8# *# i#)
-- getSB# gets a value from a byte shifting from a non-zero start
-- bit within the byte.
getSB# :: Int# -> Int#
getSB# !i# =
case i# of
0# -> let !(W8# w#) = bs `BS.index` (I# s8#)
in (word8ToInt w#) `uncheckedIShiftRL#` hop#
_ -> let !(W8# w#) = bs `BS.index` (I# (s8# +# i#))
!shft# = (8# *# i#) -# hop#
in (word8ToInt w#) `uncheckedIShiftL#` shft#
!vi# = mask# `andI#`
(case hop# of
0# -> case r8# of
1# -> getB# 0#
2# -> getB# 0# `orI#` getB# 1#
3# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2#
4# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2# `orI#` getB# 3#
5# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2# `orI#` getB# 3# `orI#`
getB# 4#
6# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2# `orI#` getB# 3# `orI#`
getB# 4# `orI#` getB# 5#
7# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2# `orI#` getB# 3# `orI#`
getB# 4# `orI#` getB# 5# `orI#`
getB# 6#
8# -> getB# 0# `orI#` getB# 1# `orI#`
getB# 2# `orI#` getB# 3# `orI#`
getB# 4# `orI#` getB# 5# `orI#`
getB# 6# `orI#` getB# 7#
-- This is the catch-all loop for other sizes
-- not addressed above.
_ -> let join !(W8# w#) !(I# a#) =
I# ((a# `uncheckedIShiftL#` 8#)
`orI#` (word8ToInt w#))
bs' = BS.take (I# (r8# +# 2#))
$ BS.drop (I# s8#) bs
!(I# v#) = BS.foldr join (0::Int) bs'
in mask# `andI#` (v# `uncheckedIShiftRL#` hop#)
_ -> case r8# of
1# -> getSB# 0#
2# -> getSB# 0# `orI#` getSB# 1#
3# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2#
4# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3#
5# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4#
6# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4# `orI#` getSB# 5#
7# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4# `orI#` getSB# 5# `orI#`
getSB# 6#
8# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4# `orI#` getSB# 5# `orI#`
getSB# 6# `orI#` getSB# 7#
-- n.b. these are hand-unrolled cases for common
-- sizes this is called for.
9# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4# `orI#` getSB# 5# `orI#`
getSB# 6# `orI#` getSB# 7# `orI#`
getSB# 8#
18# -> getSB# 0# `orI#` getSB# 1# `orI#`
getSB# 2# `orI#` getSB# 3# `orI#`
getSB# 4# `orI#` getSB# 5# `orI#`
getSB# 6# `orI#` getSB# 7# `orI#`
getSB# 8# `orI#` getSB# 9# `orI#`
getSB# 10# `orI#` getSB# 11# `orI#`
getSB# 12# `orI#` getSB# 13# `orI#`
getSB# 14# `orI#` getSB# 15# `orI#`
getSB# 16# `orI#` getSB# 17#
-- This is the catch-all loop for other sizes
-- not addressed above.
_ -> let join !(W8# w#) !(I# a#) =
I# ((a# `uncheckedIShiftL#` 8#)
`orI#` (word8ToInt w#))
bs' = BS.take (I# (r8# +# 2#))
$ BS.drop (I# s8#) bs
!(I# v#) = BS.foldr join (0::Int) bs'
in mask# `andI#` (v# `uncheckedIShiftRL#` hop#)
)
in Right $ \_ -> (# vi#, updPos# #)
else Left "Attempt to read bits past limit"
else
-- BitString stores an Int, but number of extracted bits is larger than
-- an Int can represent.
Left "Attempt to extracted large value"
-- Basic Interface -------------------------------------------------------------
-- | Read zeros up to an alignment of 32-bits.
align32bits :: GetBits ()
align32bits = GetBits $ \ !pos# inp ->
let !(# curBit#, ttlBits# #) = pos#
!s32# = curBit# `andI#` 31#
!r32# = 32# -# s32# -- num bits to reach next 32-bit boundary
nonZero = "alignments @" <> show (I# curBit#)
<> " not zeroes up to 32-bit boundary"
in if isTrue# (s32# ==# 0#)
then (# Right (), pos# #)
else case extractFromByteString ttlBits# curBit# r32# inp of
Right getRes ->
let !(# vi#, newPos# #) = getRes ()
in if isTrue# (vi# ==# 0#)
then (# Right (), (# newPos#, ttlBits# #) #)
else (# Left nonZero, pos# #)
Left e -> (# Left e, pos# #)
-- | Read out n bits as a @BitString@.
fixed :: NumBits -> GetBits BitString
fixed !(Bits' (I# n#)) = GetBits
$ \ !s@(# cur#, lim# #) ->
\inp ->
case extractFromByteString lim# cur# n# inp of
Right getRes ->
let !(# v#, p# #) = getRes ()
in (# pure $ toBitString (Bits' (I# n#)) (I# v#)
, (# p#, lim# #)
#)
Left e -> (# Left e, s #)
-- | Read out n bytes as a @ByteString@, aligning to a 32-bit boundary before and
-- after.
bytestring :: NumBytes -> GetBits ByteString
bytestring n@(Bytes' nbytes) = do
align32bits
r <- GetBits
$ \ !(# pos#, lim# #) ->
\inp ->
let !sbyte# = pos# `uncheckedIShiftRL#` 3# -- known to be aligned
!endAt# = pos# +# bitCount# (bytesToBits n)
!end# = (# endAt#, lim# #)
err = "Sub-bytestring attempted beyond end of input bytestring"
in if isTrue# (endAt# <=# lim#)
then (# pure $ BS.take nbytes $ BS.drop (I# sbyte#) inp, end# #)
else (# Left err, end# #)
align32bits
return r
-- | Add a label to the error tag stack.
label :: String -> GetBits a -> GetBits a
label l m = GetBits $ \ !pos# inp ->
let !(# j, n# #) = unGetBits m pos# inp
in case j of
Left e -> (# Left $ e <> "\n " <> l, n# #)
Right r -> (# Right r, n# #)
-- | Isolate input to a sub-span of the specified byte length.
isolate :: NumBytes -> GetBits a -> GetBits a
isolate ws m =
GetBits $ \ !(# pos#, lim# #) ->
\inp ->
let !l# = pos# +# bitCount# (bytesToBits ws)
!(# r, (# x#, _ #) #) = unGetBits m (# pos#, l# #) inp
in (# r, (# x#, lim# #) #)
-- | Try to parse something, returning Nothing when it fails.
try :: GetBits a -> GetBits (Maybe a)
try m = (Just <$> m) `mplus` return Nothing
-- | Skips the specified number of bits
skip :: NumBits -> GetBits ()
skip !(Bits' (I# n#)) =
GetBits $ \ !(# cur#, lim# #) ->
let !newLoc# = cur# +# n#
!newPos# = (# newLoc#, lim# #)
in if isTrue# (newLoc# ># lim#)
then \_ -> (# Left "skipped past end of bytestring"
, newPos#
#)
else \_ -> (# Right (), newPos# #)