ghc/compiler/utils/Binary.hs

   1 {-# OPTIONS -cpp #-}
   2 --
   3 -- (c) The University of Glasgow 2002
   4 --
   5 -- Binary I/O library, with special tweaks for GHC
   6
   7 module Binary
   8   ( {-type-}  Bin,
   9     {-class-} Binary(..),
  10     {-type-}  BinHandle,
  11
  12    openBinIO, openBinIO_,
  13    openBinMem,
  14 --   closeBin,
  15
  16    getUserData,
  17
  18    seekBin,
  19    tellBin,
  20    castBin,
  21
  22    writeBinMem,
  23    readBinMem,
  24
  25    isEOFBin,
  26
  27    -- for writing instances:
  28    putByte,
  29    getByte,
  30
  31    -- lazy Bin I/O
  32    lazyGet,
  33    lazyPut,
  34
  35    -- GHC only:
  36    ByteArray(..),
  37    getByteArray,
  38    putByteArray,
  39
  40    getBinFileWithDict,  -- :: Binary a => FilePath -> IO a
  41    putBinFileWithDict,  -- :: Binary a => FilePath -> Module -> a -> IO ()
  42
  43   ) where
  44
  45 #include "MachDeps.h"
  46
  47 import {-# SOURCE #-} Module
  48 import FastString
  49 import Unique
  50 import UniqFM
  51
  52 #if __GLASGOW_HASKELL__ < 503
  53 import IOExts
  54 import Bits
  55 import Int
  56 import Word
  57 import Char
  58 import Monad
  59 import Exception
  60 import GlaExts hiding (ByteArray, newByteArray, freezeByteArray)
  61 import Array
  62 import IO
  63 import PrelIOBase               ( IOError(..), IOErrorType(..) )
  64 import PrelReal                 ( Ratio(..) )
  65 import PrelIOBase               ( IO(..) )
  66 #else
  67 import Data.Array.IO
  68 import Data.Array
  69 import Data.Bits
  70 import Data.Int
  71 import Data.Word
  72 import Data.IORef
  73 import Data.Char                ( ord, chr )
  74 import Data.Array.Base          ( unsafeRead, unsafeWrite )
  75 import Control.Monad            ( when )
  76 import Control.Exception        ( throw )
  77 import System.IO as IO
  78 import System.IO.Unsafe         ( unsafeInterleaveIO )
  79 import System.IO.Error          ( mkIOError, eofErrorType )
  80 import GHC.Real                 ( Ratio(..) )
  81 import GHC.Exts
  82 import GHC.IOBase               ( IO(..) )
  83 import GHC.Word                 ( Word8(..) )
  84 #endif
  85
  86 #if __GLASGOW_HASKELL__ < 503
  87 type BinArray = MutableByteArray RealWorld Int
  88 newArray_ bounds     = stToIO (newCharArray bounds)
  89 unsafeWrite arr ix e = stToIO (writeWord8Array arr ix e)
  90 unsafeRead  arr ix   = stToIO (readWord8Array arr ix)
  91 newByteArray#        = newCharArray#
  92 hPutArray h arr sz   = hPutBufBAFull h arr sz
  93 hGetArray h sz       = hGetBufBAFull h sz
  94
  95 mkIOError :: IOErrorType -> String -> Maybe Handle -> Maybe FilePath -> Exception
  96 mkIOError t location maybe_hdl maybe_filename
  97   = IOException (IOError maybe_hdl t location "")
  98
  99 eofErrorType = EOF
 100
 101 #ifndef SIZEOF_HSINT
 102 #define SIZEOF_HSINT  INT_SIZE_IN_BYTES
 103 #endif
 104
 105 #ifndef SIZEOF_HSWORD
 106 #define SIZEOF_HSWORD WORD_SIZE_IN_BYTES
 107 #endif
 108
 109 #else
 110 type BinArray = IOUArray Int Word8
 111 #endif
 112
 113 data BinHandle
 114   = BinMem {            -- binary data stored in an unboxed array
 115      state :: BinHandleState,   -- sigh, need parameterized modules :-)
 116      off_r :: !FastMutInt,              -- the current offset
 117      sz_r  :: !FastMutInt,              -- size of the array (cached)
 118      arr_r :: !(IORef BinArray)         -- the array (bounds: (0,size-1))
 119     }
 120         -- XXX: should really store a "high water mark" for dumping out
 121         -- the binary data to a file.
 122
 123   | BinIO {             -- binary data stored in a file
 124      state :: BinHandleState,
 125      off_r :: !FastMutInt,              -- the current offset (cached)
 126      hdl   :: !IO.Handle                -- the file handle (must be seekable)
 127    }
 128         -- cache the file ptr in BinIO; using hTell is too expensive
 129         -- to call repeatedly.  If anyone else is modifying this Handle
 130         -- at the same time, we'll be screwed.
 131
 132 newtype Bin a = BinPtr Int
 133   deriving (Eq, Ord, Show, Bounded)
 134
 135 castBin :: Bin a -> Bin b
 136 castBin (BinPtr i) = BinPtr i
 137
 138 class Binary a where
 139     put_   :: BinHandle -> a -> IO ()
 140     put    :: BinHandle -> a -> IO (Bin a)
 141     get    :: BinHandle -> IO a
 142
 143     -- define one of put_, put.  Use of put_ is recommended because it
 144     -- is more likely that tail-calls can kick in, and we rarely need the
 145     -- position return value.
 146     put_ bh a = do put bh a; return ()
 147     put bh a  = do p <- tellBin bh; put_ bh a; return p
 148
 149 putAt  :: Binary a => BinHandle -> Bin a -> a -> IO ()
 150 putAt bh p x = do seekBin bh p; put bh x; return ()
 151
 152 getAt  :: Binary a => BinHandle -> Bin a -> IO a
 153 getAt bh p = do seekBin bh p; get bh
 154
 155 openBinIO_ :: IO.Handle -> IO BinHandle
 156 openBinIO_ h = openBinIO h noBinHandleUserData
 157
 158 openBinIO :: IO.Handle -> Module -> IO BinHandle
 159 openBinIO h mod = do
 160   r <- newFastMutInt
 161   writeFastMutInt r 0
 162   state <- newWriteState mod
 163   return (BinIO state r h)
 164
 165 openBinMem :: Int -> Module -> IO BinHandle
 166 openBinMem size mod
 167  | size <= 0 = error "Data.Binary.openBinMem: size must be >= 0"
 168  | otherwise = do
 169    arr <- newArray_ (0,size-1)
 170    arr_r <- newIORef arr
 171    ix_r <- newFastMutInt
 172    writeFastMutInt ix_r 0
 173    sz_r <- newFastMutInt
 174    writeFastMutInt sz_r size
 175    state <- newWriteState mod
 176    return (BinMem state ix_r sz_r arr_r)
 177
 178 noBinHandleUserData = error "Binary.BinHandle: no user data"
 179
 180 getUserData :: BinHandle -> BinHandleState
 181 getUserData bh = state bh
 182
 183 tellBin :: BinHandle -> IO (Bin a)
 184 tellBin (BinIO  _ r _)   = do ix <- readFastMutInt r; return (BinPtr ix)
 185 tellBin (BinMem _ r _ _) = do ix <- readFastMutInt r; return (BinPtr ix)
 186
 187 seekBin :: BinHandle -> Bin a -> IO ()
 188 seekBin (BinIO _ ix_r h) (BinPtr p) = do
 189   writeFastMutInt ix_r p
 190   hSeek h AbsoluteSeek (fromIntegral p)
 191 seekBin h@(BinMem _ ix_r sz_r a) (BinPtr p) = do
 192   sz <- readFastMutInt sz_r
 193   if (p >= sz)
 194         then do expandBin h p; writeFastMutInt ix_r p
 195         else writeFastMutInt ix_r p
 196
 197 isEOFBin :: BinHandle -> IO Bool
 198 isEOFBin (BinMem _ ix_r sz_r a) = do
 199   ix <- readFastMutInt ix_r
 200   sz <- readFastMutInt sz_r
 201   return (ix >= sz)
 202 isEOFBin (BinIO _ ix_r h) = hIsEOF h
 203
 204 writeBinMem :: BinHandle -> FilePath -> IO ()
 205 writeBinMem (BinIO _ _ _) _ = error "Data.Binary.writeBinMem: not a memory handle"
 206 writeBinMem (BinMem _ ix_r sz_r arr_r) fn = do
 207   h <- openFile fn WriteMode
 208   arr <- readIORef arr_r
 209   ix  <- readFastMutInt ix_r
 210   hPutArray h arr ix
 211 #if __GLASGOW_HASKELL__ < 500
 212   -- workaround a bug in ghc 4.08's implementation of hPutBuf (it doesn't
 213   -- set the FILEOBJ_RW_WRITTEN flag on the file object, so the file doens't
 214   -- get flushed properly).  Adding an extra '\0' doens't do any harm.
 215   hPutChar h '\0'
 216 #endif
 217   hClose h
 218
 219 readBinMem :: FilePath -> IO BinHandle
 220 readBinMem filename = do
 221   h <- openFile filename ReadMode
 222   filesize' <- hFileSize h
 223   let filesize = fromIntegral filesize'
 224   arr <- newArray_ (0,filesize-1)
 225   count <- hGetArray h arr filesize
 226   when (count /= filesize)
 227         (error ("Binary.readBinMem: only read " ++ show count ++ " bytes"))
 228   hClose h
 229   arr_r <- newIORef arr
 230   ix_r <- newFastMutInt
 231   writeFastMutInt ix_r 0
 232   sz_r <- newFastMutInt
 233   writeFastMutInt sz_r filesize
 234   return (BinMem initReadState ix_r sz_r arr_r)
 235
 236 -- expand the size of the array to include a specified offset
 237 expandBin :: BinHandle -> Int -> IO ()
 238 expandBin (BinMem _ ix_r sz_r arr_r) off = do
 239    sz <- readFastMutInt sz_r
 240    let sz' = head (dropWhile (<= off) (iterate (* 2) sz))
 241    arr <- readIORef arr_r
 242    arr' <- newArray_ (0,sz'-1)
 243    sequence_ [ unsafeRead arr i >>= unsafeWrite arr' i
 244              | i <- [ 0 .. sz-1 ] ]
 245    writeFastMutInt sz_r sz'
 246    writeIORef arr_r arr'
 247    hPutStrLn stderr ("expanding to size: " ++ show sz')
 248    return ()
 249 expandBin (BinIO _ _ _) _ = return ()
 250         -- no need to expand a file, we'll assume they expand by themselves.
 251
 252 -- -----------------------------------------------------------------------------
 253 -- Low-level reading/writing of bytes
 254
 255 putWord8 :: BinHandle -> Word8 -> IO ()
 256 putWord8 h@(BinMem _ ix_r sz_r arr_r) w = do
 257     ix <- readFastMutInt ix_r
 258     sz <- readFastMutInt sz_r
 259         -- double the size of the array if it overflows
 260     if (ix >= sz)
 261         then do expandBin h ix
 262                 putWord8 h w
 263         else do arr <- readIORef arr_r
 264                 unsafeWrite arr ix w
 265                 writeFastMutInt ix_r (ix+1)
 266                 return ()
 267 putWord8 (BinIO _ ix_r h) w = do
 268     ix <- readFastMutInt ix_r
 269     hPutChar h (chr (fromIntegral w))   -- XXX not really correct
 270     writeFastMutInt ix_r (ix+1)
 271     return ()
 272
 273 getWord8 :: BinHandle -> IO Word8
 274 getWord8 (BinMem _ ix_r sz_r arr_r) = do
 275     ix <- readFastMutInt ix_r
 276     sz <- readFastMutInt sz_r
 277     when (ix >= sz)  $
 278         throw (mkIOError eofErrorType "Data.Binary.getWord8" Nothing Nothing)
 279     arr <- readIORef arr_r
 280     w <- unsafeRead arr ix
 281     writeFastMutInt ix_r (ix+1)
 282     return w
 283 getWord8 (BinIO _ ix_r h) = do
 284     ix <- readFastMutInt ix_r
 285     c <- hGetChar h
 286     writeFastMutInt ix_r (ix+1)
 287     return (fromIntegral (ord c))       -- XXX not really correct
 288
 289 putByte :: BinHandle -> Word8 -> IO ()
 290 putByte bh w = put_ bh w
 291
 292 getByte :: BinHandle -> IO Word8
 293 getByte = getWord8
 294
 295 -- -----------------------------------------------------------------------------
 296 -- Primitve Word writes
 297
 298 instance Binary Word8 where
 299   put_ = putWord8
 300   get  = getWord8
 301
 302 instance Binary Word16 where
 303   put_ h w = do -- XXX too slow.. inline putWord8?
 304     putByte h (fromIntegral (w `shiftR` 8))
 305     putByte h (fromIntegral (w .&. 0xff))
 306   get h = do
 307     w1 <- getWord8 h
 308     w2 <- getWord8 h
 309     return ((fromIntegral w1 `shiftL` 8) .|. fromIntegral w2)
 310
 311
 312 instance Binary Word32 where
 313   put_ h w = do
 314     putByte h (fromIntegral (w `shiftR` 24))
 315     putByte h (fromIntegral ((w `shiftR` 16) .&. 0xff))
 316     putByte h (fromIntegral ((w `shiftR` 8)  .&. 0xff))
 317     putByte h (fromIntegral (w .&. 0xff))
 318   get h = do
 319     w1 <- getWord8 h
 320     w2 <- getWord8 h
 321     w3 <- getWord8 h
 322     w4 <- getWord8 h
 323     return ((fromIntegral w1 `shiftL` 24) .|.
 324             (fromIntegral w2 `shiftL` 16) .|.
 325             (fromIntegral w3 `shiftL`  8) .|.
 326             (fromIntegral w4))
 327
 328
 329 instance Binary Word64 where
 330   put_ h w = do
 331     putByte h (fromIntegral (w `shiftR` 56))
 332     putByte h (fromIntegral ((w `shiftR` 48) .&. 0xff))
 333     putByte h (fromIntegral ((w `shiftR` 40) .&. 0xff))
 334     putByte h (fromIntegral ((w `shiftR` 32) .&. 0xff))
 335     putByte h (fromIntegral ((w `shiftR` 24) .&. 0xff))
 336     putByte h (fromIntegral ((w `shiftR` 16) .&. 0xff))
 337     putByte h (fromIntegral ((w `shiftR`  8) .&. 0xff))
 338     putByte h (fromIntegral (w .&. 0xff))
 339   get h = do
 340     w1 <- getWord8 h
 341     w2 <- getWord8 h
 342     w3 <- getWord8 h
 343     w4 <- getWord8 h
 344     w5 <- getWord8 h
 345     w6 <- getWord8 h
 346     w7 <- getWord8 h
 347     w8 <- getWord8 h
 348     return ((fromIntegral w1 `shiftL` 56) .|.
 349             (fromIntegral w2 `shiftL` 48) .|.
 350             (fromIntegral w3 `shiftL` 40) .|.
 351             (fromIntegral w4 `shiftL` 32) .|.
 352             (fromIntegral w5 `shiftL` 24) .|.
 353             (fromIntegral w6 `shiftL` 16) .|.
 354             (fromIntegral w7 `shiftL`  8) .|.
 355             (fromIntegral w8))
 356
 357 -- -----------------------------------------------------------------------------
 358 -- Primitve Int writes
 359
 360 instance Binary Int8 where
 361   put_ h w = put_ h (fromIntegral w :: Word8)
 362   get h    = do w <- get h; return (fromIntegral (w::Word8))
 363
 364 instance Binary Int16 where
 365   put_ h w = put_ h (fromIntegral w :: Word16)
 366   get h    = do w <- get h; return (fromIntegral (w::Word16))
 367
 368 instance Binary Int32 where
 369   put_ h w = put_ h (fromIntegral w :: Word32)
 370   get h    = do w <- get h; return (fromIntegral (w::Word32))
 371
 372 instance Binary Int64 where
 373   put_ h w = put_ h (fromIntegral w :: Word64)
 374   get h    = do w <- get h; return (fromIntegral (w::Word64))
 375
 376 -- -----------------------------------------------------------------------------
 377 -- Instances for standard types
 378
 379 instance Binary () where
 380     put_ bh () = return ()
 381     get  _     = return ()
 382 --    getF bh p  = case getBitsF bh 0 p of (_,b) -> ((),b)
 383
 384 instance Binary Bool where
 385     put_ bh b = putByte bh (fromIntegral (fromEnum b))
 386     get  bh   = do x <- getWord8 bh; return (toEnum (fromIntegral x))
 387 --    getF bh p = case getBitsF bh 1 p of (x,b) -> (toEnum x,b)
 388
 389 instance Binary Char where
 390     put_  bh c = put_ bh (fromIntegral (ord c) :: Word32)
 391     get  bh   = do x <- get bh; return (chr (fromIntegral (x :: Word32)))
 392 --    getF bh p = case getBitsF bh 8 p of (x,b) -> (toEnum x,b)
 393
 394 instance Binary Int where
 395 #if SIZEOF_HSINT == 4
 396     put_ bh i = put_ bh (fromIntegral i :: Int32)
 397     get  bh = do
 398         x <- get bh
 399         return (fromIntegral (x :: Int32))
 400 #elif SIZEOF_HSINT == 8
 401     put_ bh i = put_ bh (fromIntegral i :: Int64)
 402     get  bh = do
 403         x <- get bh
 404         return (fromIntegral (x :: Int64))
 405 #else
 406 #error "unsupported sizeof(HsInt)"
 407 #endif
 408 --    getF bh   = getBitsF bh 32
 409
 410 instance Binary a => Binary [a] where
 411     put_ bh []     = putByte bh 0
 412     put_ bh (x:xs) = do putByte bh 1; put_ bh x; put_ bh xs
 413     get bh         = do h <- getWord8 bh
 414                         case h of
 415                           0 -> return []
 416                           _ -> do x  <- get bh
 417                                   xs <- get bh
 418                                   return (x:xs)
 419
 420 instance (Binary a, Binary b) => Binary (a,b) where
 421     put_ bh (a,b) = do put_ bh a; put_ bh b
 422     get bh        = do a <- get bh
 423                        b <- get bh
 424                        return (a,b)
 425
 426 instance (Binary a, Binary b, Binary c) => Binary (a,b,c) where
 427     put_ bh (a,b,c) = do put_ bh a; put_ bh b; put_ bh c
 428     get bh          = do a <- get bh
 429                          b <- get bh
 430                          c <- get bh
 431                          return (a,b,c)
 432
 433 instance (Binary a, Binary b, Binary c, Binary d) => Binary (a,b,c,d) where
 434     put_ bh (a,b,c,d) = do put_ bh a; put_ bh b; put_ bh c; put_ bh d
 435     get bh          = do a <- get bh
 436                          b <- get bh
 437                          c <- get bh
 438                          d <- get bh
 439                          return (a,b,c,d)
 440
 441 instance Binary a => Binary (Maybe a) where
 442     put_ bh Nothing  = putByte bh 0
 443     put_ bh (Just a) = do putByte bh 1; put_ bh a
 444     get bh           = do h <- getWord8 bh
 445                           case h of
 446                             0 -> return Nothing
 447                             _ -> do x <- get bh; return (Just x)
 448
 449 instance (Binary a, Binary b) => Binary (Either a b) where
 450     put_ bh (Left  a) = do putByte bh 0; put_ bh a
 451     put_ bh (Right b) = do putByte bh 1; put_ bh b
 452     get bh            = do h <- getWord8 bh
 453                            case h of
 454                              0 -> do a <- get bh ; return (Left a)
 455                              _ -> do b <- get bh ; return (Right b)
 456
 457 #ifdef __GLASGOW_HASKELL__
 458 instance Binary Integer where
 459     put_ bh (S# i#) = do putByte bh 0; put_ bh (I# i#)
 460     put_ bh (J# s# a#) = do
 461         p <- putByte bh 1;
 462         put_ bh (I# s#)
 463         let sz# = sizeofByteArray# a#  -- in *bytes*
 464         put_ bh (I# sz#)  -- in *bytes*
 465         putByteArray bh a# sz#
 466
 467     get bh = do
 468         b <- getByte bh
 469         case b of
 470           0 -> do (I# i#) <- get bh
 471                   return (S# i#)
 472           _ -> do (I# s#) <- get bh
 473                   sz <- get bh
 474                   (BA a#) <- getByteArray bh sz
 475                   return (J# s# a#)
 476
 477 putByteArray :: BinHandle -> ByteArray# -> Int# -> IO ()
 478 putByteArray bh a s# = loop 0#
 479   where loop n#
 480            | n# ==# s# = return ()
 481            | otherwise = do
 482                 putByte bh (indexByteArray a n#)
 483                 loop (n# +# 1#)
 484
 485 getByteArray :: BinHandle -> Int -> IO ByteArray
 486 getByteArray bh (I# sz) = do
 487   (MBA arr) <- newByteArray sz
 488   let loop n
 489            | n ==# sz = return ()
 490            | otherwise = do
 491                 w <- getByte bh
 492                 writeByteArray arr n w
 493                 loop (n +# 1#)
 494   loop 0#
 495   freezeByteArray arr
 496
 497
 498 data ByteArray = BA ByteArray#
 499 data MBA = MBA (MutableByteArray# RealWorld)
 500
 501 newByteArray :: Int# -> IO MBA
 502 newByteArray sz = IO $ \s ->
 503   case newByteArray# sz s of { (# s, arr #) ->
 504   (# s, MBA arr #) }
 505
 506 freezeByteArray :: MutableByteArray# RealWorld -> IO ByteArray
 507 freezeByteArray arr = IO $ \s ->
 508   case unsafeFreezeByteArray# arr s of { (# s, arr #) ->
 509   (# s, BA arr #) }
 510
 511 writeByteArray :: MutableByteArray# RealWorld -> Int# -> Word8 -> IO ()
 512
 513 #if __GLASGOW_HASKELL__ < 503
 514 writeByteArray arr i w8 = IO $ \s ->
 515   case word8ToWord w8 of { W# w# ->
 516   case writeCharArray# arr i (chr# (word2Int# w#)) s  of { s ->
 517   (# s , () #) }}
 518 #else
 519 writeByteArray arr i (W8# w) = IO $ \s ->
 520   case writeWord8Array# arr i w s of { s ->
 521   (# s, () #) }
 522 #endif
 523
 524 #if __GLASGOW_HASKELL__ < 503
 525 indexByteArray a# n# = fromIntegral (I# (ord# (indexCharArray# a# n#)))
 526 #else
 527 indexByteArray a# n# = W8# (indexWord8Array# a# n#)
 528 #endif
 529
 530 instance (Integral a, Binary a) => Binary (Ratio a) where
 531     put_ bh (a :% b) = do put_ bh a; put_ bh b
 532     get bh = do a <- get bh; b <- get bh; return (a :% b)
 533 #endif
 534
 535 instance Binary (Bin a) where
 536   put_ bh (BinPtr i) = put_ bh i
 537   get bh = do i <- get bh; return (BinPtr i)
 538
 539 -- -----------------------------------------------------------------------------
 540 -- unboxed mutable Ints
 541
 542 #ifdef __GLASGOW_HASKELL__
 543 data FastMutInt = FastMutInt (MutableByteArray# RealWorld)
 544
 545 newFastMutInt = IO $ \s ->
 546   case newByteArray# size s of { (# s, arr #) ->
 547   (# s, FastMutInt arr #) }
 548   where I# size = SIZEOF_HSWORD
 549
 550 readFastMutInt (FastMutInt arr) = IO $ \s ->
 551   case readIntArray# arr 0# s of { (# s, i #) ->
 552   (# s, I# i #) }
 553
 554 writeFastMutInt (FastMutInt arr) (I# i) = IO $ \s ->
 555   case writeIntArray# arr 0# i s of { s ->
 556   (# s, () #) }
 557 #endif
 558
 559 -- -----------------------------------------------------------------------------
 560 -- Lazy reading/writing
 561
 562 lazyPut :: Binary a => BinHandle -> a -> IO ()
 563 lazyPut bh a = do
 564         -- output the obj with a ptr to skip over it:
 565     pre_a <- tellBin bh
 566     put_ bh pre_a       -- save a slot for the ptr
 567     put_ bh a           -- dump the object
 568     q <- tellBin bh     -- q = ptr to after object
 569     putAt bh pre_a q    -- fill in slot before a with ptr to q
 570     seekBin bh q        -- finally carry on writing at q
 571
 572 lazyGet :: Binary a => BinHandle -> IO a
 573 lazyGet bh = do
 574     p <- get bh         -- a BinPtr
 575     p_a <- tellBin bh
 576     a <- unsafeInterleaveIO (getAt bh p_a)
 577     seekBin bh p -- skip over the object for now
 578     return a
 579
 580 -- -----------------------------------------------------------------------------
 581 -- BinHandleState
 582
 583 type BinHandleState =
 584         (Module,
 585          IORef Int,
 586          IORef (UniqFM (Int,FastString)),
 587          Array Int FastString)
 588
 589 initReadState :: BinHandleState
 590 initReadState = (undef, undef, undef, undef)
 591
 592 newWriteState :: Module -> IO BinHandleState
 593 newWriteState m = do
 594   j_r <- newIORef 0
 595   out_r <- newIORef emptyUFM
 596   return (m,j_r,out_r,undef)
 597
 598 undef = error "Binary.BinHandleState"
 599
 600 -- -----------------------------------------------------------------------------
 601 -- FastString binary interface
 602
 603 getBinFileWithDict :: Binary a => FilePath -> IO a
 604 getBinFileWithDict file_path = do
 605   bh <- Binary.readBinMem file_path
 606   dict_p <- Binary.get bh               -- get the dictionary ptr
 607   data_p <- tellBin bh
 608   seekBin bh dict_p
 609   dict <- getDictionary bh
 610   seekBin bh data_p
 611   let (mod, j_r, out_r, _) = state bh
 612   get bh{ state = (mod,j_r,out_r,dict) }
 613
 614 initBinMemSize = (1024*1024) :: Int
 615
 616 putBinFileWithDict :: Binary a => FilePath -> Module -> a -> IO ()
 617 putBinFileWithDict file_path mod a = do
 618   bh <- openBinMem initBinMemSize mod
 619   p <- tellBin bh
 620   put_ bh p             -- placeholder for ptr to dictionary
 621   put_ bh a
 622   let (_, j_r, fm_r, _) = state bh
 623   j <- readIORef j_r
 624   fm <- readIORef fm_r
 625   dict_p <- tellBin bh
 626   putAt bh p dict_p     -- fill in the placeholder
 627   seekBin bh dict_p     -- seek back to the end of the file
 628   putDictionary bh j (constructDictionary j fm)
 629   writeBinMem bh file_path
 630
 631 type Dictionary = Array Int FastString
 632         -- should be 0-indexed
 633
 634 putDictionary :: BinHandle -> Int -> Dictionary -> IO ()
 635 putDictionary bh sz dict = do
 636   put_ bh sz
 637   mapM_ (putFS bh) (elems dict)
 638
 639 getDictionary :: BinHandle -> IO Dictionary
 640 getDictionary bh = do
 641   sz <- get bh
 642   elems <- sequence (take sz (repeat (getFS bh)))
 643   return (listArray (0,sz-1) elems)
 644
 645 constructDictionary :: Int -> UniqFM (Int,FastString) -> Dictionary
 646 constructDictionary j fm = array (0,j-1) (eltsUFM fm)
 647
 648 putFS bh (FastString id l ba) = do
 649   put_ bh (I# l)
 650   putByteArray bh ba l
 651 putFS bh s = error ("Binary.put_(FastString): " ++ unpackFS s)
 652         -- Note: the length of the FastString is *not* the same as
 653         -- the size of the ByteArray: the latter is rounded up to a
 654         -- multiple of the word size.
 655
 656 getFS bh = do
 657   (I# l) <- get bh
 658   (BA ba) <- getByteArray bh (I# l)
 659   return (mkFastSubStringBA# ba 0# l)
 660         -- XXX ToDo: one too many copies here
 661
 662 instance Binary FastString where
 663   put_ bh f@(FastString id l ba) =
 664     case getUserData bh of { (_, j_r, out_r, dict) -> do
 665     out <- readIORef out_r
 666     let uniq = getUnique f
 667     case lookupUFM out uniq of
 668         Just (j,f)  -> put_ bh j
 669         Nothing -> do
 670            j <- readIORef j_r
 671            put_ bh j
 672            writeIORef j_r (j+1)
 673            writeIORef out_r (addToUFM out uniq (j,f))
 674     }
 675   put_ bh s = error ("Binary.put_(FastString): " ++ show (unpackFS s))
 676
 677   get bh = do
 678         j <- get bh
 679         case getUserData bh of (_, _, _, arr) -> return (arr ! j)