@@ -386,6 +386,7 @@ data Primitive
386386 | W64 Int Word64
387387 | BS Int B.ByteString
388388 | LBS Int L.ByteString
389+ | IsEmpty
389390 deriving (Eq, Show)
390391
391392type Program = [Primitive]
@@ -408,6 +409,7 @@ instance Arbitrary Primitive where
408409 , do n <- choose (0,10)
409410 cs <- vector n
410411 return (LBS n (L.pack cs))
412+ , return IsEmpty
411413 ]
412414 shrink p =
413415 let snk c x = map (\x' -> c (bitreq x') x') (shrinker x) in
@@ -418,6 +420,7 @@ instance Arbitrary Primitive where
418420 W64 _ x -> snk W64 x
419421 BS _ bs -> let ws = B.unpack bs in map (\ws' -> BS (length ws') (B.pack ws')) (shrink ws)
420422 LBS _ lbs -> let ws = L.unpack lbs in map (\ws' -> LBS (length ws') (L.pack ws')) (shrink ws)
423+ IsEmpty -> []
421424
422425prop_primitive :: Primitive -> Property
423426prop_primitive prim = property $
@@ -430,10 +433,10 @@ prop_primitive prim = property $
430433prop_program :: Program -> Property
431434prop_program program = property $
432435 let p = mapM_ putPrimitive program
433− g = mapM getPrimitive program
436+ g = verifyProgram (8 * fromIntegral (L.length lbs)) program
434437 lbs = runPut (runBitPut p)
435438 r = runGet (runBitGet g) lbs
436− in r == program
439+ in r
437440
438441putPrimitive :: Primitive -> BitPut ()
439442putPrimitive p =
@@ -444,6 +447,7 @@ putPrimitive p =
444447 W64 n x -> putWord64be n x
445448 BS _ bs -> putByteString bs
446449 LBS _ lbs -> mapM_ putByteString (L.toChunks lbs)
450+ IsEmpty -> return ()
447451
448452getPrimitive :: Primitive -> BitGet Primitive
449453getPrimitive p =
@@ -454,6 +458,32 @@ getPrimitive p =
454458 W64 n _ -> W64 n <$> getWord64be n
455459 BS n _ -> BS n <$> getByteString n
456460 LBS n _ -> LBS n <$> getLazyByteString n
461+ IsEmpty -> isEmpty >> return IsEmpty
462+
463+
464+verifyProgram :: Int -> Program -> BitGet Bool
465+verifyProgram totalLength ps0 = go 0 ps0
466+ where
467+ go _ [] = return True
468+ go pos (p:ps) =
469+ case p of
470+ W8 n x -> check x (getWord8 n) >> go (pos+n) ps
471+ W16 n x -> check x (getWord16be n) >> go (pos+n) ps
472+ W32 n x -> check x (getWord32be n) >> go (pos+n) ps
473+ W64 n x -> check x (getWord64be n) >> go (pos+n) ps
474+ BS n x -> check x (getByteString n) >> go (pos+(8*n)) ps
475+ LBS n x -> check x (getLazyByteString n) >> go (pos+(8*n)) ps
476+ IsEmpty -> do
477+ let expected = pos == totalLength
478+ actual <- isEmpty
479+ if expected == actual
480+ then go pos ps
481+ else error "isEmpty returned wrong value"
482+ check x g = do
483+ y <- g
484+ if x == y
485+ then return ()
486+ else error "fail fail fail!"
457487
458488{-
459489instance Random Word where