FastString.hs 23.3 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2
-- (c) The University of Glasgow, 1997-2006

3 4
{-# LANGUAGE BangPatterns, CPP, MagicHash, UnboxedTuples,
    GeneralizedNewtypeDeriving #-}
5
{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
6 7 8
-- We always optimise this, otherwise performance of a non-optimised
-- compiler is severely affected

9 10 11
-- |
-- There are two principal string types used internally by GHC:
--
Mateusz Kowalczyk's avatar
Mateusz Kowalczyk committed
12
-- ['FastString']
13
--
Mateusz Kowalczyk's avatar
Mateusz Kowalczyk committed
14 15 16 17 18
--   * A compact, hash-consed, representation of character strings.
--   * Comparison is O(1), and you can get a 'Unique.Unique' from them.
--   * Generated by 'fsLit'.
--   * Turn into 'Outputable.SDoc' with 'Outputable.ftext'.
--
Sylvain Henry's avatar
Sylvain Henry committed
19
-- ['PtrString']
Mateusz Kowalczyk's avatar
Mateusz Kowalczyk committed
20
--
Sylvain Henry's avatar
Sylvain Henry committed
21
--   * Pointer and size of a Latin-1 encoded string.
Mateusz Kowalczyk's avatar
Mateusz Kowalczyk committed
22
--   * Practically no operations.
23
--   * Outputting them is fast.
Mateusz Kowalczyk's avatar
Mateusz Kowalczyk committed
24 25
--   * Generated by 'sLit'.
--   * Turn into 'Outputable.SDoc' with 'Outputable.ptext'
26 27 28 29
--   * Requires manual memory management.
--     Improper use may lead to memory leaks or dangling pointers.
--   * It assumes Latin-1 as the encoding, therefore it cannot represent
--     arbitrary Unicode strings.
30
--
Sylvain Henry's avatar
Sylvain Henry committed
31
-- Use 'PtrString' unless you want the facilities of 'FastString'.
32 33
module FastString
       (
34
        -- * ByteString
35 36
        bytesFS,            -- :: FastString -> ByteString
        fastStringToByteString, -- = bytesFS (kept for haddock)
37
        mkFastStringByteString,
38
        fastZStringToByteString,
39
        unsafeMkByteString,
40

Ian Lynagh's avatar
Ian Lynagh committed
41 42 43 44 45 46
        -- * FastZString
        FastZString,
        hPutFZS,
        zString,
        lengthFZS,

Ian Lynagh's avatar
Ian Lynagh committed
47 48
        -- * FastStrings
        FastString(..),     -- not abstract, for now.
49

Ian Lynagh's avatar
Ian Lynagh committed
50
        -- ** Construction
51
        fsLit,
52
        mkFastString,
Ian Lynagh's avatar
Ian Lynagh committed
53
        mkFastStringBytes,
54
        mkFastStringByteList,
Ian Lynagh's avatar
Ian Lynagh committed
55 56
        mkFastStringForeignPtr,
        mkFastString#,
57

Ian Lynagh's avatar
Ian Lynagh committed
58 59
        -- ** Deconstruction
        unpackFS,           -- :: FastString -> String
60

Ian Lynagh's avatar
Ian Lynagh committed
61 62
        -- ** Encoding
        zEncodeFS,
63

Ian Lynagh's avatar
Ian Lynagh committed
64
        -- ** Operations
65
        uniqueOfFS,
Ian Lynagh's avatar
Ian Lynagh committed
66 67 68
        lengthFS,
        nullFS,
        appendFS,
69 70
        headFS,
        tailFS,
Ian Lynagh's avatar
Ian Lynagh committed
71
        concatFS,
72
        consFS,
Ian Lynagh's avatar
Ian Lynagh committed
73
        nilFS,
74
        isUnderscoreFS,
75

76
        -- ** Outputting
77 78
        hPutFS,

Ian Lynagh's avatar
Ian Lynagh committed
79 80
        -- ** Internal
        getFastStringTable,
81
        getFastStringZEncCounter,
82

Sylvain Henry's avatar
Sylvain Henry committed
83 84
        -- * PtrStrings
        PtrString (..),
Jan Stolarek's avatar
Jan Stolarek committed
85

86 87
        -- ** Construction
        sLit,
Sylvain Henry's avatar
Sylvain Henry committed
88 89
        mkPtrString#,
        mkPtrString,
Jan Stolarek's avatar
Jan Stolarek committed
90

91
        -- ** Deconstruction
Sylvain Henry's avatar
Sylvain Henry committed
92
        unpackPtrString,
Jan Stolarek's avatar
Jan Stolarek committed
93

94
        -- ** Operations
Sylvain Henry's avatar
Sylvain Henry committed
95
        lengthPS
96 97
       ) where

98
#include "HsVersions.h"
99

100 101
import GhcPrelude as Prelude

102
import Encoding
103
import FastFunctions
104
import PlainPanic
105
import Util
106

107
import Control.Concurrent.MVar
108
import Control.DeepSeq
109
import Control.Monad
110
import Data.ByteString (ByteString)
111 112 113 114
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Char8    as BSC
import qualified Data.ByteString.Internal as BS
import qualified Data.ByteString.Unsafe   as BS
115
import Foreign.C
thomie's avatar
thomie committed
116
import GHC.Exts
117
import System.IO
118
import Data.Data
119
import Data.IORef
120
import Data.Char
121
import Data.Semigroup as Semi
122

123
import GHC.IO
124

125
import Foreign
126

Ben Gamari's avatar
Ben Gamari committed
127
#if GHC_STAGE >= 2
128 129 130
import GHC.Conc.Sync    (sharedCAF)
#endif

Sylvain Henry's avatar
Sylvain Henry committed
131
import GHC.Base         ( unpackCString#, unpackNBytes# )
132 133


134 135 136 137 138
-- | Gives the UTF-8 encoded bytes corresponding to a 'FastString'
bytesFS :: FastString -> ByteString
bytesFS f = fs_bs f

{-# DEPRECATED fastStringToByteString "Use `bytesFS` instead" #-}
139
fastStringToByteString :: FastString -> ByteString
140
fastStringToByteString = bytesFS
141

142 143
fastZStringToByteString :: FastZString -> ByteString
fastZStringToByteString (FastZString bs) = bs
Ian Lynagh's avatar
Ian Lynagh committed
144

145
-- This will drop information if any character > '\xFF'
146 147
unsafeMkByteString :: String -> ByteString
unsafeMkByteString = BSC.pack
148

149 150
hashFastString :: FastString -> Int
hashFastString (FastString _ _ bs _)
151 152
    = inlinePerformIO $ BS.unsafeUseAsCStringLen bs $ \(ptr, len) ->
      return $ hashStr (castPtr ptr) len
153

Ian Lynagh's avatar
Ian Lynagh committed
154 155
-- -----------------------------------------------------------------------------

156
newtype FastZString = FastZString ByteString
157
  deriving NFData
Ian Lynagh's avatar
Ian Lynagh committed
158 159

hPutFZS :: Handle -> FastZString -> IO ()
160
hPutFZS handle (FastZString bs) = BS.hPut handle bs
Ian Lynagh's avatar
Ian Lynagh committed
161 162

zString :: FastZString -> String
163 164
zString (FastZString bs) =
    inlinePerformIO $ BS.unsafeUseAsCStringLen bs peekCAStringLen
Ian Lynagh's avatar
Ian Lynagh committed
165 166

lengthFZS :: FastZString -> Int
167
lengthFZS (FastZString bs) = BS.length bs
Ian Lynagh's avatar
Ian Lynagh committed
168

169
mkFastZStringString :: String -> FastZString
170
mkFastZStringString str = FastZString (BSC.pack str)
171

Ian Lynagh's avatar
Ian Lynagh committed
172
-- -----------------------------------------------------------------------------
173

174 175 176
{-| A 'FastString' is a UTF-8 encoded string together with a unique ID. All
'FastString's are stored in a global hashtable to support fast O(1)
comparison.
177

178 179
It is also associated with a lazy reference to the Z-encoding
of this string which is used by the compiler internally.
180 181
-}
data FastString = FastString {
Ian Lynagh's avatar
Ian Lynagh committed
182 183
      uniq    :: {-# UNPACK #-} !Int, -- unique id
      n_chars :: {-# UNPACK #-} !Int, -- number of chars
184
      fs_bs   :: {-# UNPACK #-} !ByteString,
185 186 187 188 189
      fs_zenc :: FastZString
      -- ^ Lazily computed z-encoding of this string.
      --
      -- Since 'FastString's are globally memoized this is computed at most
      -- once for any given string.
190
  }
191 192 193

instance Eq FastString where
  f1 == f2  =  uniq f1 == uniq f2
194

195
instance Ord FastString where
Ian Lynagh's avatar
Ian Lynagh committed
196
    -- Compares lexicographically, not by unique
197
    a <= b = case cmpFS a b of { LT -> True;  EQ -> True;  GT -> False }
Ian Lynagh's avatar
Ian Lynagh committed
198
    a <  b = case cmpFS a b of { LT -> True;  EQ -> False; GT -> False }
199
    a >= b = case cmpFS a b of { LT -> False; EQ -> True;  GT -> True  }
Ian Lynagh's avatar
Ian Lynagh committed
200 201 202 203 204
    a >  b = case cmpFS a b of { LT -> False; EQ -> False; GT -> True  }
    max x y | x >= y    =  x
            | otherwise =  y
    min x y | x <= y    =  x
            | otherwise =  y
205
    compare a b = cmpFS a b
206

207 208 209
instance IsString FastString where
    fromString = fsLit

210 211 212
instance Semi.Semigroup FastString where
    (<>) = appendFS

213 214
instance Monoid FastString where
    mempty = nilFS
215
    mappend = (Semi.<>)
216
    mconcat = concatFS
217

218 219 220
instance Show FastString where
   show fs = show (unpackFS fs)

221 222 223 224 225 226
instance Data FastString where
  -- don't traverse?
  toConstr _   = abstractConstr "FastString"
  gunfold _ _  = error "gunfold"
  dataTypeOf _ = mkNoRepType "FastString"

227 228 229
instance NFData FastString where
  rnf fs = seq fs ()

230
cmpFS :: FastString -> FastString -> Ordering
231
cmpFS f1@(FastString u1 _ _ _) f2@(FastString u2 _ _ _) =
232
  if u1 == u2 then EQ else
233
  compare (bytesFS f1) (bytesFS f2)
234

Sylvain Henry's avatar
Sylvain Henry committed
235
foreign import ccall unsafe "memcmp"
236
  memcmp :: Ptr a -> Ptr b -> Int -> IO Int
237

238 239
-- -----------------------------------------------------------------------------
-- Construction
240

241
{-
242 243 244 245 246 247
Internally, the compiler will maintain a fast string symbol table, providing
sharing and fast comparison. Creation of new @FastString@s then covertly does a
lookup, re-using the @FastString@ if there was a hit.

The design of the FastString hash table allows for lockless concurrent reads
and updates to multiple buckets with low synchronization overhead.
248

249 250
See Note [Updating the FastString table] on how it's updated.
-}
251 252
data FastStringTable = FastStringTable
  {-# UNPACK #-} !(IORef Int) -- the unique ID counter shared with all buckets
253
  {-# UNPACK #-} !(IORef Int) -- number of computed z-encodings for all buckets
254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324
  (Array# (IORef FastStringTableSegment)) -- concurrent segments

data FastStringTableSegment = FastStringTableSegment
  {-# UNPACK #-} !(MVar ()) -- the lock for write in each segment
  {-# UNPACK #-} !(IORef Int) -- the number of elements
  (MutableArray# RealWorld [FastString]) -- buckets in this segment

{-
Following parameters are determined based on:

* Benchmark based on testsuite/tests/utils/should_run/T14854.hs
* Stats of @echo :browse | ghc --interactive -dfaststring-stats >/dev/null@:
  on 2018-10-24, we have 13920 entries.
-}
segmentBits, numSegments, segmentMask, initialNumBuckets :: Int
segmentBits = 8
numSegments = 256   -- bit segmentBits
segmentMask = 0xff  -- bit segmentBits - 1
initialNumBuckets = 64

hashToSegment# :: Int# -> Int#
hashToSegment# hash# = hash# `andI#` segmentMask#
  where
    !(I# segmentMask#) = segmentMask

hashToIndex# :: MutableArray# RealWorld [FastString] -> Int# -> Int#
hashToIndex# buckets# hash# =
  (hash# `uncheckedIShiftRL#` segmentBits#) `remInt#` size#
  where
    !(I# segmentBits#) = segmentBits
    size# = sizeofMutableArray# buckets#

maybeResizeSegment :: IORef FastStringTableSegment -> IO FastStringTableSegment
maybeResizeSegment segmentRef = do
  segment@(FastStringTableSegment lock counter old#) <- readIORef segmentRef
  let oldSize# = sizeofMutableArray# old#
      newSize# = oldSize# *# 2#
  (I# n#) <- readIORef counter
  if isTrue# (n# <# newSize#) -- maximum load of 1
  then return segment
  else do
    resizedSegment@(FastStringTableSegment _ _ new#) <- IO $ \s1# ->
      case newArray# newSize# [] s1# of
        (# s2#, arr# #) -> (# s2#, FastStringTableSegment lock counter arr# #)
    forM_ [0 .. (I# oldSize#) - 1] $ \(I# i#) -> do
      fsList <- IO $ readArray# old# i#
      forM_ fsList $ \fs -> do
        let -- Shall we store in hash value in FastString instead?
            !(I# hash#) = hashFastString fs
            idx# = hashToIndex# new# hash#
        IO $ \s1# ->
          case readArray# new# idx# s1# of
            (# s2#, bucket #) -> case writeArray# new# idx# (fs: bucket) s2# of
              s3# -> (# s3#, () #)
    writeIORef segmentRef resizedSegment
    return resizedSegment

{-# NOINLINE stringTable #-}
stringTable :: FastStringTable
stringTable = unsafePerformIO $ do
  let !(I# numSegments#) = numSegments
      !(I# initialNumBuckets#) = initialNumBuckets
      loop a# i# s1#
        | isTrue# (i# ==# numSegments#) = s1#
        | otherwise = case newMVar () `unIO` s1# of
            (# s2#, lock #) -> case newIORef 0 `unIO` s2# of
              (# s3#, counter #) -> case newArray# initialNumBuckets# [] s3# of
                (# s4#, buckets# #) -> case newIORef
                    (FastStringTableSegment lock counter buckets#) `unIO` s4# of
                  (# s5#, segment #) -> case writeArray# a# i# segment s5# of
                    s6# -> loop a# (i# +# 1#) s6#
325
  uid <- newIORef 603979776 -- ord '$' * 0x01000000
326
  n_zencs <- newIORef 0
327 328 329 330
  tab <- IO $ \s1# ->
    case newArray# numSegments# (panic "string_table") s1# of
      (# s2#, arr# #) -> case loop arr# 0# s2# of
        s3# -> case unsafeFreezeArray# arr# s3# of
331 332
          (# s4#, segments# #) ->
            (# s4#, FastStringTable uid n_zencs segments# #)
333

334 335
  -- use the support wired into the RTS to share this CAF among all images of
  -- libHSghc
Ben Gamari's avatar
Ben Gamari committed
336
#if GHC_STAGE < 2
337
  return tab
338
#else
339
  sharedCAF tab getOrSetLibHSghcFastStringTable
340

Ben Gamari's avatar
Ben Gamari committed
341
-- from the RTS; thus we cannot use this mechanism when GHC_STAGE<2; the previous
342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376
-- RTS might not have this symbol
foreign import ccall unsafe "getOrSetLibHSghcFastStringTable"
  getOrSetLibHSghcFastStringTable :: Ptr a -> IO (Ptr a)
#endif

{-

We include the FastString table in the `sharedCAF` mechanism because we'd like
FastStrings created by a Core plugin to have the same uniques as corresponding
strings created by the host compiler itself.  For example, this allows plugins
to lookup known names (eg `mkTcOcc "MySpecialType"`) in the GlobalRdrEnv or
even re-invoke the parser.

In particular, the following little sanity test was failing in a plugin
prototyping safe newtype-coercions: GHC.NT.Type.NT was imported, but could not
be looked up /by the plugin/.

   let rdrName = mkModuleName "GHC.NT.Type" `mkRdrQual` mkTcOcc "NT"
   putMsgS $ showSDoc dflags $ ppr $ lookupGRE_RdrName rdrName $ mg_rdr_env guts

`mkTcOcc` involves the lookup (or creation) of a FastString.  Since the
plugin's FastString.string_table is empty, constructing the RdrName also
allocates new uniques for the FastStrings "GHC.NT.Type" and "NT".  These
uniques are almost certainly unequal to the ones that the host compiler
originally assigned to those FastStrings.  Thus the lookup fails since the
domain of the GlobalRdrEnv is affected by the RdrName's OccName's FastString's
unique.

Maintaining synchronization of the two instances of this global is rather
difficult because of the uses of `unsafePerformIO` in this module.  Not
synchronizing them risks breaking the rather major invariant that two
FastStrings with the same unique have the same string. Thus we use the
lower-level `sharedCAF` mechanism that relies on Globals.c.

-}
377

378
mkFastString# :: Addr# -> FastString
379
mkFastString# a# = mkFastStringBytes ptr (ptrStrLength ptr)
380
  where ptr = Ptr a#
381

382 383
{- Note [Updating the FastString table]

384 385 386 387 388
We use a concurrent hashtable which contains multiple segments, each hash value
always maps to the same segment. Read is lock-free, write to the a segment
should acquire a lock for that segment to avoid race condition, writes to
different segments are independent.

389 390
The procedure goes like this:

391 392 393 394 395
1. Find out which segment to operate on based on the hash value
2. Read the relevant bucket and perform a look up of the string.
3. If it exists, return it.
4. Otherwise grab a unique ID, create a new FastString and atomically attempt
   to update the relevant segment with this FastString:
396

397 398
   * Resize the segment by doubling the number of buckets when the number of
     FastStrings in this segment grows beyond the threshold.
399 400 401 402 403 404 405
   * Double check that the string is not in the bucket. Another thread may have
     inserted it while we were creating our string.
   * Return the existing FastString if it exists. The one we preemptively
     created will get GCed.
   * Otherwise, insert and return the string we created.
-}

406 407
mkFastStringWith
    :: (Int -> IORef Int-> IO FastString) -> Ptr Word8 -> Int -> IO FastString
408
mkFastStringWith mk_fs !ptr !len = do
409 410 411 412 413 414 415
  FastStringTableSegment lock _ buckets# <- readIORef segmentRef
  let idx# = hashToIndex# buckets# hash#
  bucket <- IO $ readArray# buckets# idx#
  res <- bucket_match bucket len ptr
  case res of
    Just found -> return found
    Nothing -> do
416 417 418
      -- The withMVar below is not dupable. It can lead to deadlock if it is
      -- only run partially and putMVar is not called after takeMVar.
      noDuplicate
419
      n <- get_uid
420
      new_fs <- mk_fs n n_zencs
421
      withMVar lock $ \_ -> insert new_fs
422
  where
423
    !(FastStringTable uid n_zencs segments#) = stringTable
424
    get_uid = atomicModifyIORef' uid $ \n -> (n+1,n)
425

426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454
    !(I# hash#) = hashStr ptr len
    (# segmentRef #) = indexArray# segments# (hashToSegment# hash#)
    insert fs = do
      FastStringTableSegment _ counter buckets# <- maybeResizeSegment segmentRef
      let idx# = hashToIndex# buckets# hash#
      bucket <- IO $ readArray# buckets# idx#
      res <- bucket_match bucket len ptr
      case res of
        -- The FastString was added by another thread after previous read and
        -- before we acquired the write lock.
        Just found -> return found
        Nothing -> do
          IO $ \s1# ->
            case writeArray# buckets# idx# (fs: bucket) s1# of
              s2# -> (# s2#, () #)
          modifyIORef' counter succ
          return fs

bucket_match :: [FastString] -> Int -> Ptr Word8 -> IO (Maybe FastString)
bucket_match [] _ _ = return Nothing
bucket_match (v@(FastString _ _ bs _):ls) len ptr
      | len == BS.length bs = do
         b <- BS.unsafeUseAsCString bs $ \buf ->
             cmpStringPrefix ptr (castPtr buf) len
         if b then return (Just v)
              else bucket_match ls len ptr
      | otherwise =
         bucket_match ls len ptr

455
mkFastStringBytes :: Ptr Word8 -> Int -> FastString
456 457 458 459 460
mkFastStringBytes !ptr !len =
    -- NB: Might as well use unsafeDupablePerformIO, since mkFastStringWith is
    -- idempotent.
    unsafeDupablePerformIO $
        mkFastStringWith (copyNewFastString ptr len) ptr len
461 462 463 464 465

-- | Create a 'FastString' from an existing 'ForeignPtr'; the difference
-- between this and 'mkFastStringBytes' is that we don't have to copy
-- the bytes if the string is new to the table.
mkFastStringForeignPtr :: Ptr Word8 -> ForeignPtr Word8 -> Int -> IO FastString
466 467
mkFastStringForeignPtr ptr !fp len
    = mkFastStringWith (mkNewFastString fp ptr len) ptr len
468

469 470 471
-- | Create a 'FastString' from an existing 'ForeignPtr'; the difference
-- between this and 'mkFastStringBytes' is that we don't have to copy
-- the bytes if the string is new to the table.
472 473 474 475 476 477
mkFastStringByteString :: ByteString -> FastString
mkFastStringByteString bs =
    inlinePerformIO $
      BS.unsafeUseAsCStringLen bs $ \(ptr, len) -> do
        let ptr' = castPtr ptr
        mkFastStringWith (mkNewFastStringByteString bs ptr' len) ptr' len
478

479 480
-- | Creates a UTF-8 encoded 'FastString' from a 'String'
mkFastString :: String -> FastString
Ian Lynagh's avatar
Ian Lynagh committed
481
mkFastString str =
482 483 484 485 486
  inlinePerformIO $ do
    let l = utf8EncodedLength str
    buf <- mallocForeignPtrBytes l
    withForeignPtr buf $ \ptr -> do
      utf8EncodeString ptr str
Ian Lynagh's avatar
Ian Lynagh committed
487
      mkFastStringForeignPtr ptr buf l
488

489 490
-- | Creates a 'FastString' from a UTF-8 encoded @[Word8]@
mkFastStringByteList :: [Word8] -> FastString
491
mkFastStringByteList str = mkFastStringByteString (BS.pack str)
492

493 494 495 496 497 498
-- | Creates a (lazy) Z-encoded 'FastString' from a 'String' and account
-- the number of forced z-strings into the passed 'IORef'.
mkZFastString :: IORef Int -> ByteString -> FastZString
mkZFastString n_zencs bs = unsafePerformIO $ do
  atomicModifyIORef' n_zencs $ \n -> (n+1, ())
  return $ mkFastZStringString (zEncodeString (utf8DecodeByteString bs))
499

500
mkNewFastString :: ForeignPtr Word8 -> Ptr Word8 -> Int -> Int
501 502 503 504
                -> IORef Int -> IO FastString
mkNewFastString fp ptr len uid n_zencs = do
  let bs = BS.fromForeignPtr fp 0 len
      zstr = mkZFastString n_zencs bs
505
  n_chars <- countUTF8Chars ptr len
506
  return (FastString uid n_chars bs zstr)
507

508
mkNewFastStringByteString :: ByteString -> Ptr Word8 -> Int -> Int
509 510 511
                          -> IORef Int -> IO FastString
mkNewFastStringByteString bs ptr len uid n_zencs = do
  let zstr = mkZFastString n_zencs bs
512
  n_chars <- countUTF8Chars ptr len
513
  return (FastString uid n_chars bs zstr)
514

515 516
copyNewFastString :: Ptr Word8 -> Int -> Int -> IORef Int -> IO FastString
copyNewFastString ptr len uid n_zencs = do
517
  fp <- copyBytesToForeignPtr ptr len
518 519
  let bs = BS.fromForeignPtr fp 0 len
      zstr = mkZFastString n_zencs bs
520
  n_chars <- countUTF8Chars ptr len
521
  return (FastString uid n_chars bs zstr)
522

523
copyBytesToForeignPtr :: Ptr Word8 -> Int -> IO (ForeignPtr Word8)
524 525 526 527 528
copyBytesToForeignPtr ptr len = do
  fp <- mallocForeignPtrBytes len
  withForeignPtr fp $ \ptr' -> copyBytes ptr' ptr len
  return fp

529 530 531
cmpStringPrefix :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
cmpStringPrefix ptr1 ptr2 len =
 do r <- memcmp ptr1 ptr2 len
532 533 534 535
    return (r == 0)


hashStr  :: Ptr Word8 -> Int -> Int
536
 -- use the Addr to produce a hash value between 0 & m (inclusive)
537
hashStr (Ptr a#) (I# len#) = loop 0# 0#
Ian Lynagh's avatar
Ian Lynagh committed
538
   where
thomie's avatar
thomie committed
539 540
    loop h n | isTrue# (n ==# len#) = I# h
             | otherwise  = loop h2 (n +# 1#)
541 542 543
          where
            !c = ord# (indexCharOffAddr# a# n)
            !h2 = (h *# 16777619#) `xorI#` c
544

545 546
-- -----------------------------------------------------------------------------
-- Operations
547

548 549 550
-- | Returns the length of the 'FastString' in characters
lengthFS :: FastString -> Int
lengthFS f = n_chars f
551

552
-- | Returns @True@ if the 'FastString' is empty
553
nullFS :: FastString -> Bool
554
nullFS f = BS.null (fs_bs f)
555

556
-- | Unpacks and decodes the FastString
557
unpackFS :: FastString -> String
558
unpackFS (FastString _ _ bs _) = utf8DecodeByteString bs
559

560
-- | Returns a Z-encoded version of a 'FastString'.  This might be the
561 562 563 564
-- original, if it was already Z-encoded.  The first time this
-- function is applied to a particular 'FastString', the results are
-- memoized.
--
Ian Lynagh's avatar
Ian Lynagh committed
565
zEncodeFS :: FastString -> FastZString
566
zEncodeFS (FastString _ _ _ ref) = ref
567 568

appendFS :: FastString -> FastString -> FastString
569
appendFS fs1 fs2 = mkFastStringByteString
570
                 $ BS.append (bytesFS fs1) (bytesFS fs2)
571 572

concatFS :: [FastString] -> FastString
573
concatFS = mkFastStringByteString . BS.concat . map fs_bs
574 575

headFS :: FastString -> Char
576
headFS (FastString _ 0 _ _) = panic "headFS: Empty FastString"
577 578 579
headFS (FastString _ _ bs _) =
  inlinePerformIO $ BS.unsafeUseAsCString bs $ \ptr ->
         return (fst (utf8DecodeChar (castPtr ptr)))
580 581

tailFS :: FastString -> FastString
582
tailFS (FastString _ 0 _ _) = panic "tailFS: Empty FastString"
583 584
tailFS (FastString _ _ bs _) =
    inlinePerformIO $ BS.unsafeUseAsCString bs $ \ptr ->
585
    do let (_, n) = utf8DecodeChar (castPtr ptr)
586
       return $! mkFastStringByteString (BS.drop n bs)
587 588 589 590

consFS :: Char -> FastString -> FastString
consFS c fs = mkFastString (c : unpackFS fs)

591 592
uniqueOfFS :: FastString -> Int
uniqueOfFS (FastString u _ _ _) = u
593

594
nilFS :: FastString
595
nilFS = mkFastString ""
596

597 598 599
isUnderscoreFS :: FastString -> Bool
isUnderscoreFS fs = fs == fsLit "_"

600 601 602
-- -----------------------------------------------------------------------------
-- Stats

603 604 605 606 607 608 609 610 611
getFastStringTable :: IO [[[FastString]]]
getFastStringTable =
  forM [0 .. numSegments - 1] $ \(I# i#) -> do
    let (# segmentRef #) = indexArray# segments# i#
    FastStringTableSegment _ _ buckets# <- readIORef segmentRef
    let bucketSize = I# (sizeofMutableArray# buckets#)
    forM [0 .. bucketSize - 1] $ \(I# j#) ->
      IO $ readArray# buckets# j#
  where
612 613 614 615 616 617
    !(FastStringTable _ _ segments#) = stringTable

getFastStringZEncCounter :: IO Int
getFastStringZEncCounter = readIORef n_zencs
  where
    !(FastStringTable _ n_zencs _) = stringTable
618

619 620 621
-- -----------------------------------------------------------------------------
-- Outputting 'FastString's

622 623
-- |Outputs a 'FastString' with /no decoding at all/, that is, you
-- get the actual bytes in the 'FastString' written to the 'Handle'.
624
hPutFS :: Handle -> FastString -> IO ()
625
hPutFS handle fs = BS.hPut handle $ bytesFS fs
626

627 628
-- ToDo: we'll probably want an hPutFSLocal, or something, to output
-- in the current locale's encoding (for error messages and suchlike).
629

630
-- -----------------------------------------------------------------------------
Sylvain Henry's avatar
Sylvain Henry committed
631
-- PtrStrings, here for convenience only.
632

Sylvain Henry's avatar
Sylvain Henry committed
633 634
-- | A 'PtrString' is a pointer to some array of Latin-1 encoded chars.
data PtrString = PtrString !(Ptr Word8) !Int
635

Sylvain Henry's avatar
Sylvain Henry committed
636 637 638
-- | Wrap an unboxed address into a 'PtrString'.
mkPtrString# :: Addr# -> PtrString
mkPtrString# a# = PtrString (Ptr a#) (ptrStrLength (Ptr a#))
639

Sylvain Henry's avatar
Sylvain Henry committed
640
-- | Encode a 'String' into a newly allocated 'PtrString' using Latin-1
641 642
-- encoding.  The original string must not contain non-Latin-1 characters
-- (above codepoint @0xff@).
Sylvain Henry's avatar
Sylvain Henry committed
643 644 645
{-# INLINE mkPtrString #-}
mkPtrString :: String -> PtrString
mkPtrString s =
Sylvain Henry's avatar
Sylvain Henry committed
646 647
 -- we don't use `unsafeDupablePerformIO` here to avoid potential memory leaks
 -- and because someone might be using `eqAddr#` to check for string equality.
648
 unsafePerformIO (do
Sylvain Henry's avatar
Sylvain Henry committed
649 650
   let len = length s
   p <- mallocBytes len
651 652
   let
     loop :: Int -> String -> IO ()
Sylvain Henry's avatar
Sylvain Henry committed
653
     loop !_ []    = return ()
654 655 656 657
     loop n (c:cs) = do
        pokeByteOff p n (fromIntegral (ord c) :: Word8)
        loop (1+n) cs
   loop 0 s
Sylvain Henry's avatar
Sylvain Henry committed
658
   return (PtrString p len)
659 660
 )

Sylvain Henry's avatar
Sylvain Henry committed
661 662 663 664
-- | Decode a 'PtrString' back into a 'String' using Latin-1 encoding.
-- This does not free the memory associated with 'PtrString'.
unpackPtrString :: PtrString -> String
unpackPtrString (PtrString (Ptr p#) (I# n#)) = unpackNBytes# p# n#
665

Sylvain Henry's avatar
Sylvain Henry committed
666 667 668
-- | Return the length of a 'PtrString'
lengthPS :: PtrString -> Int
lengthPS (PtrString _ n) = n
669

670 671 672
-- -----------------------------------------------------------------------------
-- under the carpet

Sylvain Henry's avatar
Sylvain Henry committed
673
foreign import ccall unsafe "strlen"
674
  ptrStrLength :: Ptr Word8 -> Int
675

Ian Lynagh's avatar
Ian Lynagh committed
676
{-# NOINLINE sLit #-}
Sylvain Henry's avatar
Sylvain Henry committed
677 678
sLit :: String -> PtrString
sLit x  = mkPtrString x
Ian Lynagh's avatar
Ian Lynagh committed
679 680 681 682 683 684

{-# NOINLINE fsLit #-}
fsLit :: String -> FastString
fsLit x = mkFastString x

{-# RULES "slit"
Sylvain Henry's avatar
Sylvain Henry committed
685
    forall x . sLit  (unpackCString# x) = mkPtrString#  x #-}
Ian Lynagh's avatar
Ian Lynagh committed
686 687
{-# RULES "fslit"
    forall x . fsLit (unpackCString# x) = mkFastString# x #-}