compiler/utils/Binary.hs

   1 {-# OPTIONS -cpp #-}
   2 --
   3 -- (c) The University of Glasgow 2002-2006
   4 --
   5 -- Binary I/O library, with special tweaks for GHC
   6 --
   7 -- Based on the nhc98 Binary library, which is copyright
   8 -- (c) Malcolm Wallace and Colin Runciman, University of York, 1998.
   9 -- Under the terms of the license for that software, we must tell you
  10 -- where you can obtain the original version of the Binary library, namely
  11 --     http://www.cs.york.ac.uk/fp/nhc98/
  12
  13 module Binary
  14   ( {-type-}  Bin,
  15     {-class-} Binary(..),
  16     {-type-}  BinHandle,
  17
  18    openBinIO, openBinIO_,
  19    openBinMem,
  20 --   closeBin,
  21
  22    seekBin,
  23    tellBin,
  24    castBin,
  25
  26    writeBinMem,
  27    readBinMem,
  28
  29    isEOFBin,
  30
  31    putAt, getAt,
  32
  33    -- for writing instances:
  34    putByte,
  35    getByte,
  36
  37    -- lazy Bin I/O
  38    lazyGet,
  39    lazyPut,
  40
  41 #ifdef __GLASGOW_HASKELL__
  42    -- GHC only:
  43    ByteArray(..),
  44    getByteArray,
  45    putByteArray,
  46 #endif
  47
  48    UserData(..), getUserData, setUserData,
  49    newReadState, newWriteState,
  50    putDictionary, getDictionary,
  51   ) where
  52
  53 #include "HsVersions.h"
  54
  55 -- The *host* architecture version:
  56 #include "MachDeps.h"
  57
  58 import {-# SOURCE #-} Name (Name)
  59 import FastString
  60 import Unique
  61 import Panic
  62 import UniqFM
  63 import FastMutInt
  64
  65 import Foreign
  66 import Data.Array.IO
  67 import Data.Array
  68 import Data.Bits
  69 import Data.Int
  70 import Data.Word
  71 import Data.IORef
  72 import Data.Char                ( ord, chr )
  73 import Data.Array.Base          ( unsafeRead, unsafeWrite )
  74 import Control.Monad            ( when )
  75 import System.IO as IO
  76 import System.IO.Unsafe         ( unsafeInterleaveIO )
  77 import System.IO.Error          ( mkIOError, eofErrorType )
  78 import GHC.Real                 ( Ratio(..) )
  79 import GHC.Exts
  80 import GHC.IOBase               ( IO(..) )
  81 import GHC.Word                 ( Word8(..) )
  82 #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 601
  83 -- openFileEx is available from the lang package, but we want to
  84 -- be independent of hslibs libraries.
  85 import GHC.Handle               ( openFileEx, IOModeEx(..) )
  86 #else
  87 import System.IO                ( openBinaryFile )
  88 #endif
  89
  90 #if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 601
  91 openBinaryFile f mode = openFileEx f (BinaryMode mode)
  92 #endif
  93
  94 type BinArray = IOUArray Int Word8
  95
  96 ---------------------------------------------------------------
  97 -- BinHandle
  98 ---------------------------------------------------------------
  99
 100 data BinHandle
 101   = BinMem {                     -- binary data stored in an unboxed array
 102      bh_usr :: UserData,         -- sigh, need parameterized modules :-)
 103      _off_r :: !FastMutInt,      -- the current offset
 104      _sz_r  :: !FastMutInt,      -- size of the array (cached)
 105      _arr_r :: !(IORef BinArray) -- the array (bounds: (0,size-1))
 106     }
 107         -- XXX: should really store a "high water mark" for dumping out
 108         -- the binary data to a file.
 109
 110   | BinIO {                     -- binary data stored in a file
 111      bh_usr :: UserData,
 112      _off_r :: !FastMutInt,     -- the current offset (cached)
 113      _hdl   :: !IO.Handle       -- the file handle (must be seekable)
 114    }
 115         -- cache the file ptr in BinIO; using hTell is too expensive
 116         -- to call repeatedly.  If anyone else is modifying this Handle
 117         -- at the same time, we'll be screwed.
 118
 119 getUserData :: BinHandle -> UserData
 120 getUserData bh = bh_usr bh
 121
 122 setUserData :: BinHandle -> UserData -> BinHandle
 123 setUserData bh us = bh { bh_usr = us }
 124
 125
 126 ---------------------------------------------------------------
 127 -- Bin
 128 ---------------------------------------------------------------
 129
 130 newtype Bin a = BinPtr Int
 131   deriving (Eq, Ord, Show, Bounded)
 132
 133 castBin :: Bin a -> Bin b
 134 castBin (BinPtr i) = BinPtr i
 135
 136 ---------------------------------------------------------------
 137 -- class Binary
 138 ---------------------------------------------------------------
 139
 140 class Binary a where
 141     put_   :: BinHandle -> a -> IO ()
 142     put    :: BinHandle -> a -> IO (Bin a)
 143     get    :: BinHandle -> IO a
 144
 145     -- define one of put_, put.  Use of put_ is recommended because it
 146     -- is more likely that tail-calls can kick in, and we rarely need the
 147     -- position return value.
 148     put_ bh a = do put bh a; return ()
 149     put bh a  = do p <- tellBin bh; put_ bh a; return p
 150
 151 putAt  :: Binary a => BinHandle -> Bin a -> a -> IO ()
 152 putAt bh p x = do seekBin bh p; put bh x; return ()
 153
 154 getAt  :: Binary a => BinHandle -> Bin a -> IO a
 155 getAt bh p = do seekBin bh p; get bh
 156
 157 openBinIO_ :: IO.Handle -> IO BinHandle
 158 openBinIO_ h = openBinIO h
 159
 160 openBinIO :: IO.Handle -> IO BinHandle
 161 openBinIO h = do
 162   r <- newFastMutInt
 163   writeFastMutInt r 0
 164   return (BinIO noUserData r h)
 165
 166 openBinMem :: Int -> IO BinHandle
 167 openBinMem size
 168  | size <= 0 = error "Data.Binary.openBinMem: size must be >= 0"
 169  | otherwise = do
 170    arr <- newArray_ (0,size-1)
 171    arr_r <- newIORef arr
 172    ix_r <- newFastMutInt
 173    writeFastMutInt ix_r 0
 174    sz_r <- newFastMutInt
 175    writeFastMutInt sz_r size
 176    return (BinMem noUserData ix_r sz_r arr_r)
 177
 178 tellBin :: BinHandle -> IO (Bin a)
 179 tellBin (BinIO  _ r _)   = do ix <- readFastMutInt r; return (BinPtr ix)
 180 tellBin (BinMem _ r _ _) = do ix <- readFastMutInt r; return (BinPtr ix)
 181
 182 seekBin :: BinHandle -> Bin a -> IO ()
 183 seekBin (BinIO _ ix_r h) (BinPtr p) = do
 184   writeFastMutInt ix_r p
 185   hSeek h AbsoluteSeek (fromIntegral p)
 186 seekBin h@(BinMem _ ix_r sz_r _) (BinPtr p) = do
 187   sz <- readFastMutInt sz_r
 188   if (p >= sz)
 189         then do expandBin h p; writeFastMutInt ix_r p
 190         else writeFastMutInt ix_r p
 191
 192 isEOFBin :: BinHandle -> IO Bool
 193 isEOFBin (BinMem _ ix_r sz_r _) = do
 194   ix <- readFastMutInt ix_r
 195   sz <- readFastMutInt sz_r
 196   return (ix >= sz)
 197 isEOFBin (BinIO _ _ h) = hIsEOF h
 198
 199 writeBinMem :: BinHandle -> FilePath -> IO ()
 200 writeBinMem (BinIO _ _ _) _ = error "Data.Binary.writeBinMem: not a memory handle"
 201 writeBinMem (BinMem _ ix_r _ arr_r) fn = do
 202   h <- openBinaryFile fn WriteMode
 203   arr <- readIORef arr_r
 204   ix  <- readFastMutInt ix_r
 205   hPutArray h arr ix
 206   hClose h
 207
 208 readBinMem :: FilePath -> IO BinHandle
 209 -- Return a BinHandle with a totally undefined State
 210 readBinMem filename = do
 211   h <- openBinaryFile filename ReadMode
 212   filesize' <- hFileSize h
 213   let filesize = fromIntegral filesize'
 214   arr <- newArray_ (0,filesize-1)
 215   count <- hGetArray h arr filesize
 216   when (count /= filesize)
 217        (error ("Binary.readBinMem: only read " ++ show count ++ " bytes"))
 218   hClose h
 219   arr_r <- newIORef arr
 220   ix_r <- newFastMutInt
 221   writeFastMutInt ix_r 0
 222   sz_r <- newFastMutInt
 223   writeFastMutInt sz_r filesize
 224   return (BinMem noUserData ix_r sz_r arr_r)
 225
 226 -- expand the size of the array to include a specified offset
 227 expandBin :: BinHandle -> Int -> IO ()
 228 expandBin (BinMem _ _ sz_r arr_r) off = do
 229    sz <- readFastMutInt sz_r
 230    let sz' = head (dropWhile (<= off) (iterate (* 2) sz))
 231    arr <- readIORef arr_r
 232    arr' <- newArray_ (0,sz'-1)
 233    sequence_ [ unsafeRead arr i >>= unsafeWrite arr' i
 234              | i <- [ 0 .. sz-1 ] ]
 235    writeFastMutInt sz_r sz'
 236    writeIORef arr_r arr'
 237 #ifdef DEBUG
 238    hPutStrLn stderr ("Binary: expanding to size: " ++ show sz')
 239 #endif
 240    return ()
 241 expandBin (BinIO _ _ _) _ = return ()
 242 -- no need to expand a file, we'll assume they expand by themselves.
 243
 244 -- -----------------------------------------------------------------------------
 245 -- Low-level reading/writing of bytes
 246
 247 putWord8 :: BinHandle -> Word8 -> IO ()
 248 putWord8 h@(BinMem _ ix_r sz_r arr_r) w = do
 249     ix <- readFastMutInt ix_r
 250     sz <- readFastMutInt sz_r
 251     -- double the size of the array if it overflows
 252     if (ix >= sz)
 253         then do expandBin h ix
 254                 putWord8 h w
 255         else do arr <- readIORef arr_r
 256                 unsafeWrite arr ix w
 257                 writeFastMutInt ix_r (ix+1)
 258                 return ()
 259 putWord8 (BinIO _ ix_r h) w = do
 260     ix <- readFastMutInt ix_r
 261     hPutChar h (chr (fromIntegral w)) -- XXX not really correct
 262     writeFastMutInt ix_r (ix+1)
 263     return ()
 264
 265 getWord8 :: BinHandle -> IO Word8
 266 getWord8 (BinMem _ ix_r sz_r arr_r) = do
 267     ix <- readFastMutInt ix_r
 268     sz <- readFastMutInt sz_r
 269     when (ix >= sz) $
 270         ioError (mkIOError eofErrorType "Data.Binary.getWord8" Nothing Nothing)
 271     arr <- readIORef arr_r
 272     w <- unsafeRead arr ix
 273     writeFastMutInt ix_r (ix+1)
 274     return w
 275 getWord8 (BinIO _ ix_r h) = do
 276     ix <- readFastMutInt ix_r
 277     c <- hGetChar h
 278     writeFastMutInt ix_r (ix+1)
 279     return $! (fromIntegral (ord c)) -- XXX not really correct
 280
 281 putByte :: BinHandle -> Word8 -> IO ()
 282 putByte bh w = put_ bh w
 283
 284 getByte :: BinHandle -> IO Word8
 285 getByte = getWord8
 286
 287 -- -----------------------------------------------------------------------------
 288 -- Primitve Word writes
 289
 290 instance Binary Word8 where
 291   put_ = putWord8
 292   get  = getWord8
 293
 294 instance Binary Word16 where
 295   put_ h w = do -- XXX too slow.. inline putWord8?
 296     putByte h (fromIntegral (w `shiftR` 8))
 297     putByte h (fromIntegral (w .&. 0xff))
 298   get h = do
 299     w1 <- getWord8 h
 300     w2 <- getWord8 h
 301     return $! ((fromIntegral w1 `shiftL` 8) .|. fromIntegral w2)
 302
 303
 304 instance Binary Word32 where
 305   put_ h w = do
 306     putByte h (fromIntegral (w `shiftR` 24))
 307     putByte h (fromIntegral ((w `shiftR` 16) .&. 0xff))
 308     putByte h (fromIntegral ((w `shiftR` 8)  .&. 0xff))
 309     putByte h (fromIntegral (w .&. 0xff))
 310   get h = do
 311     w1 <- getWord8 h
 312     w2 <- getWord8 h
 313     w3 <- getWord8 h
 314     w4 <- getWord8 h
 315     return $! ((fromIntegral w1 `shiftL` 24) .|.
 316                (fromIntegral w2 `shiftL` 16) .|.
 317                (fromIntegral w3 `shiftL`  8) .|.
 318                (fromIntegral w4))
 319
 320 instance Binary Word64 where
 321   put_ h w = do
 322     putByte h (fromIntegral (w `shiftR` 56))
 323     putByte h (fromIntegral ((w `shiftR` 48) .&. 0xff))
 324     putByte h (fromIntegral ((w `shiftR` 40) .&. 0xff))
 325     putByte h (fromIntegral ((w `shiftR` 32) .&. 0xff))
 326     putByte h (fromIntegral ((w `shiftR` 24) .&. 0xff))
 327     putByte h (fromIntegral ((w `shiftR` 16) .&. 0xff))
 328     putByte h (fromIntegral ((w `shiftR`  8) .&. 0xff))
 329     putByte h (fromIntegral (w .&. 0xff))
 330   get h = do
 331     w1 <- getWord8 h
 332     w2 <- getWord8 h
 333     w3 <- getWord8 h
 334     w4 <- getWord8 h
 335     w5 <- getWord8 h
 336     w6 <- getWord8 h
 337     w7 <- getWord8 h
 338     w8 <- getWord8 h
 339     return $! ((fromIntegral w1 `shiftL` 56) .|.
 340                (fromIntegral w2 `shiftL` 48) .|.
 341                (fromIntegral w3 `shiftL` 40) .|.
 342                (fromIntegral w4 `shiftL` 32) .|.
 343                (fromIntegral w5 `shiftL` 24) .|.
 344                (fromIntegral w6 `shiftL` 16) .|.
 345                (fromIntegral w7 `shiftL`  8) .|.
 346                (fromIntegral w8))
 347
 348 -- -----------------------------------------------------------------------------
 349 -- Primitve Int writes
 350
 351 instance Binary Int8 where
 352   put_ h w = put_ h (fromIntegral w :: Word8)
 353   get h    = do w <- get h; return $! (fromIntegral (w::Word8))
 354
 355 instance Binary Int16 where
 356   put_ h w = put_ h (fromIntegral w :: Word16)
 357   get h    = do w <- get h; return $! (fromIntegral (w::Word16))
 358
 359 instance Binary Int32 where
 360   put_ h w = put_ h (fromIntegral w :: Word32)
 361   get h    = do w <- get h; return $! (fromIntegral (w::Word32))
 362
 363 instance Binary Int64 where
 364   put_ h w = put_ h (fromIntegral w :: Word64)
 365   get h    = do w <- get h; return $! (fromIntegral (w::Word64))
 366
 367 -- -----------------------------------------------------------------------------
 368 -- Instances for standard types
 369
 370 instance Binary () where
 371     put_ _ () = return ()
 372     get  _    = return ()
 373 --    getF bh p  = case getBitsF bh 0 p of (_,b) -> ((),b)
 374
 375 instance Binary Bool where
 376     put_ bh b = putByte bh (fromIntegral (fromEnum b))
 377     get  bh   = do x <- getWord8 bh; return $! (toEnum (fromIntegral x))
 378 --    getF bh p = case getBitsF bh 1 p of (x,b) -> (toEnum x,b)
 379
 380 instance Binary Char where
 381     put_  bh c = put_ bh (fromIntegral (ord c) :: Word32)
 382     get  bh   = do x <- get bh; return $! (chr (fromIntegral (x :: Word32)))
 383 --    getF bh p = case getBitsF bh 8 p of (x,b) -> (toEnum x,b)
 384
 385 instance Binary Int where
 386 #if SIZEOF_HSINT == 4
 387     put_ bh i = put_ bh (fromIntegral i :: Int32)
 388     get  bh = do
 389         x <- get bh
 390         return $! (fromIntegral (x :: Int32))
 391 #elif SIZEOF_HSINT == 8
 392     put_ bh i = put_ bh (fromIntegral i :: Int64)
 393     get  bh = do
 394         x <- get bh
 395         return $! (fromIntegral (x :: Int64))
 396 #else
 397 #error "unsupported sizeof(HsInt)"
 398 #endif
 399 --    getF bh   = getBitsF bh 32
 400
 401 instance Binary a => Binary [a] where
 402     put_ bh l = do
 403         let len = length l
 404         if (len < 0xff)
 405           then putByte bh (fromIntegral len :: Word8)
 406           else do putByte bh 0xff; put_ bh (fromIntegral len :: Word32)
 407         mapM_ (put_ bh) l
 408     get bh = do
 409         b <- getByte bh
 410         len <- if b == 0xff
 411                   then get bh
 412                   else return (fromIntegral b :: Word32)
 413         let loop 0 = return []
 414             loop n = do a <- get bh; as <- loop (n-1); return (a:as)
 415         loop len
 416
 417 instance (Binary a, Binary b) => Binary (a,b) where
 418     put_ bh (a,b) = do put_ bh a; put_ bh b
 419     get bh        = do a <- get bh
 420                        b <- get bh
 421                        return (a,b)
 422
 423 instance (Binary a, Binary b, Binary c) => Binary (a,b,c) where
 424     put_ bh (a,b,c) = do put_ bh a; put_ bh b; put_ bh c
 425     get bh          = do a <- get bh
 426                          b <- get bh
 427                          c <- get bh
 428                          return (a,b,c)
 429
 430 instance (Binary a, Binary b, Binary c, Binary d) => Binary (a,b,c,d) where
 431     put_ bh (a,b,c,d) = do put_ bh a; put_ bh b; put_ bh c; put_ bh d
 432     get bh          = do a <- get bh
 433                          b <- get bh
 434                          c <- get bh
 435                          d <- get bh
 436                          return (a,b,c,d)
 437
 438 instance Binary a => Binary (Maybe a) where
 439     put_ bh Nothing  = putByte bh 0
 440     put_ bh (Just a) = do putByte bh 1; put_ bh a
 441     get bh           = do h <- getWord8 bh
 442                           case h of
 443                             0 -> return Nothing
 444                             _ -> do x <- get bh; return (Just x)
 445
 446 instance (Binary a, Binary b) => Binary (Either a b) where
 447     put_ bh (Left  a) = do putByte bh 0; put_ bh a
 448     put_ bh (Right b) = do putByte bh 1; put_ bh b
 449     get bh            = do h <- getWord8 bh
 450                            case h of
 451                              0 -> do a <- get bh ; return (Left a)
 452                              _ -> do b <- get bh ; return (Right b)
 453
 454 #if defined(__GLASGOW_HASKELL__) || 1
 455 --to quote binary-0.3 on this code idea,
 456 --
 457 -- TODO  This instance is not architecture portable.  GMP stores numbers as
 458 -- arrays of machine sized words, so the byte format is not portable across
 459 -- architectures with different endianess and word size.
 460 --
 461 -- This makes it hard (impossible) to make an equivalent instance
 462 -- with code that is compilable with non-GHC.  Do we need any instance
 463 -- Binary Integer, and if so, does it have to be blazing fast?  Or can
 464 -- we just change this instance to be portable like the rest of the
 465 -- instances? (binary package has code to steal for that)
 466 --
 467 -- yes, we need Binary Integer and Binary Rational in basicTypes/Literal.lhs
 468
 469 instance Binary Integer where
 470     put_ bh (S# i#) = do putByte bh 0; put_ bh (I# i#)
 471     put_ bh (J# s# a#) = do
 472         putByte bh 1
 473         put_ bh (I# s#)
 474         let sz# = sizeofByteArray# a#  -- in *bytes*
 475         put_ bh (I# sz#)  -- in *bytes*
 476         putByteArray bh a# sz#
 477
 478     get bh = do
 479         b <- getByte bh
 480         case b of
 481           0 -> do (I# i#) <- get bh
 482                   return (S# i#)
 483           _ -> do (I# s#) <- get bh
 484                   sz <- get bh
 485                   (BA a#) <- getByteArray bh sz
 486                   return (J# s# a#)
 487
 488 -- As for the rest of this code, even though this module
 489 -- exports it, it doesn't seem to be used anywhere else
 490 -- in GHC!
 491
 492 putByteArray :: BinHandle -> ByteArray# -> Int# -> IO ()
 493 putByteArray bh a s# = loop 0#
 494   where loop n#
 495            | n# ==# s# = return ()
 496            | otherwise = do
 497                 putByte bh (indexByteArray a n#)
 498                 loop (n# +# 1#)
 499
 500 getByteArray :: BinHandle -> Int -> IO ByteArray
 501 getByteArray bh (I# sz) = do
 502   (MBA arr) <- newByteArray sz
 503   let loop n
 504            | n ==# sz = return ()
 505            | otherwise = do
 506                 w <- getByte bh
 507                 writeByteArray arr n w
 508                 loop (n +# 1#)
 509   loop 0#
 510   freezeByteArray arr
 511
 512
 513 data ByteArray = BA ByteArray#
 514 data MBA = MBA (MutableByteArray# RealWorld)
 515
 516 newByteArray :: Int# -> IO MBA
 517 newByteArray sz = IO $ \s ->
 518   case newByteArray# sz s of { (# s, arr #) ->
 519   (# s, MBA arr #) }
 520
 521 freezeByteArray :: MutableByteArray# RealWorld -> IO ByteArray
 522 freezeByteArray arr = IO $ \s ->
 523   case unsafeFreezeByteArray# arr s of { (# s, arr #) ->
 524   (# s, BA arr #) }
 525
 526 writeByteArray :: MutableByteArray# RealWorld -> Int# -> Word8 -> IO ()
 527 writeByteArray arr i (W8# w) = IO $ \s ->
 528   case writeWord8Array# arr i w s of { s ->
 529   (# s, () #) }
 530
 531 indexByteArray :: ByteArray# -> Int# -> Word8
 532 indexByteArray a# n# = W8# (indexWord8Array# a# n#)
 533
 534 instance (Integral a, Binary a) => Binary (Ratio a) where
 535     put_ bh (a :% b) = do put_ bh a; put_ bh b
 536     get bh = do a <- get bh; b <- get bh; return (a :% b)
 537 #endif
 538
 539 instance Binary (Bin a) where
 540   put_ bh (BinPtr i) = put_ bh i
 541   get bh = do i <- get bh; return (BinPtr i)
 542
 543 -- -----------------------------------------------------------------------------
 544 -- Lazy reading/writing
 545
 546 lazyPut :: Binary a => BinHandle -> a -> IO ()
 547 lazyPut bh a = do
 548     -- output the obj with a ptr to skip over it:
 549     pre_a <- tellBin bh
 550     put_ bh pre_a       -- save a slot for the ptr
 551     put_ bh a           -- dump the object
 552     q <- tellBin bh     -- q = ptr to after object
 553     putAt bh pre_a q    -- fill in slot before a with ptr to q
 554     seekBin bh q        -- finally carry on writing at q
 555
 556 lazyGet :: Binary a => BinHandle -> IO a
 557 lazyGet bh = do
 558     p <- get bh -- a BinPtr
 559     p_a <- tellBin bh
 560     a <- unsafeInterleaveIO (getAt bh p_a)
 561     seekBin bh p -- skip over the object for now
 562     return a
 563
 564 -- -----------------------------------------------------------------------------
 565 -- UserData
 566 -- -----------------------------------------------------------------------------
 567
 568 data UserData =
 569    UserData {
 570         -- for *deserialising* only:
 571         ud_dict   :: Dictionary,
 572         ud_symtab :: SymbolTable,
 573
 574         -- for *serialising* only:
 575         ud_dict_next :: !FastMutInt, -- The next index to use
 576         ud_dict_map  :: !(IORef (UniqFM (Int,FastString))),
 577                                 -- indexed by FastString
 578
 579         ud_symtab_next :: !FastMutInt, -- The next index to use
 580         ud_symtab_map  :: !(IORef (UniqFM (Int,Name)))
 581                                 -- indexed by Name
 582    }
 583
 584 newReadState :: Dictionary -> IO UserData
 585 newReadState dict = do
 586   dict_next <- newFastMutInt
 587   dict_map <- newIORef (undef "dict_map")
 588   symtab_next <- newFastMutInt
 589   symtab_map <- newIORef (undef "symtab_map")
 590   return UserData { ud_dict = dict,
 591                     ud_symtab = undef "symtab",
 592                     ud_dict_next = dict_next,
 593                     ud_dict_map = dict_map,
 594                     ud_symtab_next = symtab_next,
 595                     ud_symtab_map = symtab_map
 596                    }
 597
 598 newWriteState :: IO UserData
 599 newWriteState = do
 600   dict_next <- newFastMutInt
 601   writeFastMutInt dict_next 0
 602   dict_map <- newIORef emptyUFM
 603   symtab_next <- newFastMutInt
 604   writeFastMutInt symtab_next 0
 605   symtab_map <- newIORef emptyUFM
 606   return UserData { ud_dict = undef "dict",
 607                     ud_symtab = undef "symtab",
 608                     ud_dict_next = dict_next,
 609                     ud_dict_map = dict_map,
 610                     ud_symtab_next = symtab_next,
 611                     ud_symtab_map = symtab_map
 612                    }
 613
 614 noUserData :: a
 615 noUserData = undef "UserData"
 616
 617 undef :: String -> a
 618 undef s = panic ("Binary.UserData: no " ++ s)
 619
 620 ---------------------------------------------------------
 621 -- The Dictionary
 622 ---------------------------------------------------------
 623
 624 type Dictionary = Array Int FastString -- The dictionary
 625                                        -- Should be 0-indexed
 626
 627 putDictionary :: BinHandle -> Int -> UniqFM (Int,FastString) -> IO ()
 628 putDictionary bh sz dict = do
 629   put_ bh sz
 630   mapM_ (putFS bh) (elems (array (0,sz-1) (eltsUFM dict)))
 631
 632 getDictionary :: BinHandle -> IO Dictionary
 633 getDictionary bh = do
 634   sz <- get bh
 635   elems <- sequence (take sz (repeat (getFS bh)))
 636   return (listArray (0,sz-1) elems)
 637
 638 ---------------------------------------------------------
 639 -- The Symbol Table
 640 ---------------------------------------------------------
 641
 642 -- On disk, the symbol table is an array of IfaceExtName, when
 643 -- reading it in we turn it into a SymbolTable.
 644
 645 type SymbolTable = Array Int Name
 646
 647 ---------------------------------------------------------
 648 -- Reading and writing FastStrings
 649 ---------------------------------------------------------
 650
 651 putFS :: BinHandle -> FastString -> IO ()
 652 putFS bh (FastString _ l _ buf _) = do
 653   put_ bh l
 654   withForeignPtr buf $ \ptr ->
 655     let
 656         go n | n == l    = return ()
 657              | otherwise = do
 658                 b <- peekElemOff ptr n
 659                 putByte bh b
 660                 go (n+1)
 661    in
 662    go 0
 663
 664 {- -- possible faster version, not quite there yet:
 665 getFS bh@BinMem{} = do
 666   (I# l) <- get bh
 667   arr <- readIORef (arr_r bh)
 668   off <- readFastMutInt (off_r bh)
 669   return $! (mkFastSubStringBA# arr off l)
 670 -}
 671 getFS :: BinHandle -> IO FastString
 672 getFS bh = do
 673   l <- get bh
 674   fp <- mallocForeignPtrBytes l
 675   withForeignPtr fp $ \ptr -> do
 676   let
 677         go n | n == l = mkFastStringForeignPtr ptr fp l
 678              | otherwise = do
 679                 b <- getByte bh
 680                 pokeElemOff ptr n b
 681                 go (n+1)
 682   --
 683   go 0
 684
 685 instance Binary FastString where
 686   put_ bh f =
 687     case getUserData bh of {
 688         UserData { ud_dict_next = j_r,
 689                    ud_dict_map = out_r} -> do
 690     out <- readIORef out_r
 691     let uniq = getUnique f
 692     case lookupUFM out uniq of
 693         Just (j, _)  -> put_ bh j
 694         Nothing -> do
 695            j <- readFastMutInt j_r
 696            put_ bh j
 697            writeFastMutInt j_r (j + 1)
 698            writeIORef out_r $! addToUFM out uniq (j, f)
 699     }
 700
 701   get bh = do
 702         j <- get bh
 703         return $! (ud_dict (getUserData bh) ! j)