Skip to content

Commit 3e4d369

Browse files
angermanangerman
committed
Make buffer adjustable
This doesn't work correctly yet :-( We also will have an issue with overrun at the end, as we grab the _words, but don't know on which byte they end. However if all trailing bytes are zero, this might be less of a problem.
1 parent c9818de commit 3e4d369

File tree

1 file changed

+53
-39
lines changed

1 file changed

+53
-39
lines changed

‎src/Data/BitCode/Writer/Monad.hs‎

Lines changed: 53 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ import Data.Foldable
2929
import Control.Monad.Trans.State.Strict
3030

3131
import qualified Data.List as L
32-
import qualified Data.ByteString as B
32+
import qualified Data.ByteString.Lazy as B
33+
import qualified Data.Binary.Put as Bin (runPut)
34+
import qualified Data.Binary as Bin (put)
3335
import Data.Sequence (Seq)
3436
import qualified Data.Sequence as Seq
3537
-- import Data.Monoid ((<>))
@@ -43,9 +45,21 @@ import Debug.Trace
4345
-- | The position in the stream.
4446
type Position = Int
4547

48+
bSize :: Int
49+
-- ensure the order is correct
50+
bToOrder :: BType -> BType
51+
52+
-- Word64
53+
-- type BType = Word64
54+
-- bSize = 64
55+
-- bToOrder = byteSwap64
56+
57+
type BType = Word8
58+
bSize = 8
59+
bToOrder = id -- Word8
4660
-- | A @Word8@ buffer, tracking the number of bits.
4761
-- I don't think those Unpacks are necessary, -funbox-small-strict-fields is on by default
48-
data Buff = Buff !Int !Word8 deriving (Eq, Ord)
62+
data Buff = Buff !Int !BType deriving (Eq, Ord)
4963

5064
mask :: (FiniteBits a, Num a) => Int -> a -> a
5165
mask n w = m .&. w
@@ -80,37 +94,37 @@ nullBuff = Buff 0 0
8094
-- In the spill case B = (5, 0b00010101) and C = (4, 0b00001111) we expect to get
8195
-- (Just 0b11110101, (1, b000000001))
8296
--
83-
addBuff :: Buff -> Buff -> (Maybe Word8, Buff)
84-
addBuff (Buff n w ) (Buff n' w' ) | n+n' < 8 = (Nothing
85-
, Buff (n+n') (w .|. (shift w' n)))
97+
addBuff :: Buff -> Buff -> (Maybe BType, Buff)
98+
addBuff (Buff n w ) (Buff n' w' ) | n+n' < bSize = (Nothing
99+
, Buff (n+n') (w .|. (shift w' n)))
86100
| otherwise = (Just (w .|. (shift w' n))
87-
, Buff ((n+n') `mod` 8) (shift w' (n-8)))
101+
, Buff ((n+n') `mod` bSize) (shift w' (n-bSize)))
88102

89103

90104
-- | Smart constructor for @Buff@. Ensures that
91105
-- the stored byte is masked properly.
92-
mkBuff :: Int -> Word8 -> Buff
106+
mkBuff :: Int -> BType -> Buff
93107
mkBuff n w = Buff n (mask n w)
94108

95109
-- | A stream is a number of Words, a buffer and a position (length of the stream) marker.
96110
data Stream f a = S
97-
{ _words :: !(f Word8)
111+
{ _words :: !(f BType)
98112
, _buffer :: !Buff
99113
, _len :: !Position
100114
}
101115

102-
deriving instance Eq (f Word8) => Eq (Stream f a)
103-
deriving instance Ord (f Word8) => Ord (Stream f a)
116+
deriving instance Eq (f BType) => Eq (Stream f a)
117+
deriving instance Ord (f BType) => Ord (Stream f a)
104118

105119
data Streams f a = Streams
106120
{ _substreams :: !(Seq (Stream f a))
107121
, _total_len :: !Position
108122
}
109123

110-
deriving instance Eq (f Word8) => Eq (Streams f a)
111-
deriving instance Ord (f Word8) => Ord (Streams f a)
124+
deriving instance Eq (f BType) => Eq (Streams f a)
125+
deriving instance Ord (f BType) => Ord (Streams f a)
112126

113-
instance ( Semigroup (f Word8)
127+
instance ( Semigroup (f BType)
114128
, Foldable f
115129
, Traversable f
116130
, Applicative f) => Semigroup (Stream f a) where
@@ -132,18 +146,18 @@ instance ( Semigroup (f Word8)
132146
(Just w''', b'') -> S (w <> w'' <> pure w''') b'' (p + p')
133147
(Nothing, b'') -> S (w <> w'') b'' (p + p')
134148
where go' :: Int -- ^ shift
135-
-> Word8 -- ^ buff
136-
-> Word8 -- ^ input
137-
-> ( Word8 -- ^ new buff
138-
, Word8 ) -- ^ output
139-
go' !n !b !w = (shift w (n-8), b .|. shift w n)
149+
-> BType -- ^ buff
150+
-> BType -- ^ input
151+
-> ( BType -- ^ new buff
152+
, BType ) -- ^ output
153+
go' !n !b !w = (shift w (n-bSize), b .|. shift w n)
140154
in r
141155

142156
{-# SPECIALIZE instance Semigroup (Stream Seq a) #-}
143157
{-# SPECIALIZE instance Semigroup (Stream [] a) #-}
144158

145-
instance ( Semigroup (f Word8)
146-
, Monoid (f Word8)
159+
instance ( Semigroup (f BType)
160+
, Monoid (f BType)
147161
, Foldable f
148162
, Traversable f
149163
, Applicative f) => Monoid (Stream f a) where
@@ -166,11 +180,11 @@ instance ( Semigroup (f Word8)
166180
(Just w''', b'') -> S (w <> w'' <> pure w''') b'' (p + p')
167181
(Nothing, b'') -> S (w <> w'') b'' (p + p')
168182
where go' :: Int -- ^ shift
169-
-> Word8 -- ^ buff
170-
-> Word8 -- ^ input
171-
-> ( Word8 -- ^ new buff
172-
, Word8 ) -- ^ output
173-
go' !n !b !w = (shift w (n-8), b .|. shift w n)
183+
-> BType -- ^ buff
184+
-> BType -- ^ input
185+
-> ( BType -- ^ new buff
186+
, BType ) -- ^ output
187+
go' !n !b !w = (shift w (n-bSize), b .|. shift w n)
174188
in r
175189

176190
{-# SPECIALIZE instance Monoid (Stream Seq a) #-}
@@ -192,8 +206,8 @@ instance Monoid (Streams f a) where
192206
{-# SPECIALIZE instance Monoid (Streams [] a) #-}
193207

194208
-- mappend is not cheap here.
195-
type ListStream = Stream [] Word8
196-
type SeqStreams = Streams Seq Word8
209+
type ListStream = Stream [] BType
210+
type SeqStreams = Streams Seq BType
197211

198212
toListStream :: Foldable f => Stream f a -> Stream [] a
199213
toListStream (S w b p) = S (toList w) b p
@@ -221,21 +235,21 @@ value (_,_,v) = v
221235

222236
runBitstream :: Position -> Bitstream a -> (ListStream, Position, a)
223237
runBitstream p (Bitstream f) = case runState f (BitstreamState mempty 0) of (a, BitstreamState ss p) -> (toListStream . runStreams $ ss, p, a)
224-
execBitstream :: Position -> Bitstream a -> [Word8]
238+
execBitstream :: Position -> Bitstream a -> [BType]
225239
execBitstream p a = _words . stream . runBitstream p $ a >> alignWord8
226240
evalBitstream :: Position -> Bitstream a -> a
227241
evalBitstream p = value . runBitstream p
228242

229-
streams :: Foldable f => f Word8 -> Buff -> Position -> SeqStreams
243+
streams :: Foldable f => f BType -> Buff -> Position -> SeqStreams
230244
streams w b p
231245
| p == 0 = mempty
232246
| otherwise = Streams (pure $ S (Seq.fromList . toList $ w) b p) p
233247

234-
{-# SPECIALIZE streams :: [Word8] -> Buff -> Position -> SeqStreams #-}
235-
bitstream :: Foldable f => f Word8 -> Buff -> Int -> Bitstream ()
248+
{-# SPECIALIZE streams :: [BType] -> Buff -> Position -> SeqStreams #-}
249+
bitstream :: Foldable f => f BType -> Buff -> Int -> Bitstream ()
236250
bitstream w b p = Bitstream $ modify' $ \(BitstreamState ss p') -> BitstreamState (ss <> streams w b p) (p + p')
237251

238-
{-# SPECIALIZE bitstream :: [Word8] -> Buff -> Int -> Bitstream () #-}
252+
{-# SPECIALIZE bitstream :: [BType] -> Buff -> Int -> Bitstream () #-}
239253
-- Monadic Bitstream API
240254

241255
withOffset :: Int -> Bitstream a -> Bitstream a
@@ -258,7 +272,7 @@ emitBit :: Bool -> Bitstream ()
258272
emitBit True = bitstream [] (Buff 1 0b00000001) 1
259273
emitBit False = bitstream [] (Buff 1 0b00000000) 1
260274

261-
emitBits :: Int -> Word8 -> Bitstream ()
275+
emitBits :: Int -> BType -> Bitstream ()
262276
emitBits 0 _ = pure ()
263277
emitBits n b | n < 8 = do
264278
-- traceM $ "emitting " ++ show n ++ " bits; value = " ++ show b
@@ -268,7 +282,7 @@ emitBits n b | n < 8 = do
268282
| otherwise = error $ "cannot emit " ++ show n ++ " bits from Word8."
269283

270284
emitWord8 :: Word8 -> Bitstream ()
271-
emitWord8 w = bitstream [w] nullBuff 8
285+
emitWord8 w = bitstream [fromIntegral w] nullBuff 8
272286

273287
emitWord32R :: Word32 -> Bitstream ()
274288
emitWord32R w = bitstream [fromIntegral (shift w (-24))
@@ -317,7 +331,7 @@ emitFixed n w | n < 8 = bitstream [] (mkBuff n' (off 0 w)) n'
317331
nullBuff
318332
64
319333
| otherwise = error $ "invalid number of bits. Cannot emit " ++ show n ++ " bits from Word64."
320-
where off :: Int -> Word64 -> Word8
334+
where off :: Int -> Word64 -> BType
321335
off n w = fromIntegral (shift w (-n))
322336
n' = fromIntegral n
323337

@@ -357,8 +371,8 @@ alignWord32 = flip mod 32 <$> loc >>= \case
357371
writeFile
358372
:: HasCallStack
359373
=> FilePath -> Bitstream a -> IO ()
360-
writeFile f = B.writeFile f . B.pack . execBitstream 0
361-
374+
writeFile f = B.writeFile f . encode' . execBitstream 0
375+
where encode' = Bin.runPut . mapM_ (Bin.put . bToOrder)
362376
-- * BitCode Header
363377
-- | put the BitCodeHeader, on darwin a special wrapper is
364378
-- apparently only required, to make it compatible with
@@ -404,8 +418,8 @@ withHeader isDarwin body = mdo
404418

405419
-- Show instances. These make parsing debug output much easier.
406420

407-
showWord8 :: Word8 -> String
408-
showWord8 w = '0':'b':(map f $ [testBit w i | i <- [0..7]])
421+
showWord8 :: BType -> String
422+
showWord8 w = '0':'b':(map f $ [testBit w i | i <- [0..bSize-1]])
409423
where f True = '1'
410424
f False = '0'
411425

0 commit comments

Comments
 (0)