{-# LANGUAGE BangPatterns, MagicHash, UnboxedTuples #-}
{-# OPTIONS_GHC -O2 #-}
-- We always optimise this, otherwise performance of a non-optimised
-- compiler is severely affected

-- -----------------------------------------------------------------------------
--
-- (c) The University of Glasgow, 1997-2006
--
-- Character encodings
--
-- -----------------------------------------------------------------------------

module Encoding (
        -- * UTF-8
        utf8DecodeChar#,
        utf8PrevChar,
        utf8CharStart,
        utf8DecodeChar,
        utf8DecodeByteString,
        utf8DecodeStringLazy,
        utf8EncodeChar,
        utf8EncodeString,
        utf8EncodedLength,
        countUTF8Chars,

        -- * Z-encoding
        zEncodeString,
        zDecodeString,

        -- * Base62-encoding
        toBase62,
        toBase62Padded
  ) where

import GhcPrelude

import Foreign
import Foreign.ForeignPtr.Unsafe
import Data.Char
import qualified Data.Char as Char
import Numeric
import GHC.IO

import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BS

import GHC.Exts

-- -----------------------------------------------------------------------------
-- UTF-8

-- We can't write the decoder as efficiently as we'd like without
-- resorting to unboxed extensions, unfortunately.  I tried to write
-- an IO version of this function, but GHC can't eliminate boxed
-- results from an IO-returning function.
--
-- We assume we can ignore overflow when parsing a multibyte character here.
-- To make this safe, we add extra sentinel bytes to unparsed UTF-8 sequences
-- before decoding them (see StringBuffer.hs).

{-# INLINE utf8DecodeChar# #-}
utf8DecodeChar# :: Addr# -> (# Char#, Int# #)
utf8DecodeChar# a# =
  let !ch0 = word2Int# (indexWord8OffAddr# a# 0#) in
  case () of
    _ | isTrue# (ch0 <=# 0x7F#) -> (# chr# ch0, 1# #)

      | isTrue# ((ch0 >=# 0xC0#) `andI#` (ch0 <=# 0xDF#)) ->
        let !ch1 = word2Int# (indexWord8OffAddr# a# 1#) in
        if isTrue# ((ch1 <# 0x80#) `orI#` (ch1 >=# 0xC0#)) then fail 1# else
        (# chr# (((ch0 -# 0xC0#) `uncheckedIShiftL#` 6#) +#
                  (ch1 -# 0x80#)),
           2# #)

      | isTrue# ((ch0 >=# 0xE0#) `andI#` (ch0 <=# 0xEF#)) ->
        let !ch1 = word2Int# (indexWord8OffAddr# a# 1#) in
        if isTrue# ((ch1 <# 0x80#) `orI#` (ch1 >=# 0xC0#)) then fail 1# else
        let !ch2 = word2Int# (indexWord8OffAddr# a# 2#) in
        if isTrue# ((ch2 <# 0x80#) `orI#` (ch2 >=# 0xC0#)) then fail 2# else
        (# chr# (((ch0 -# 0xE0#) `uncheckedIShiftL#` 12#) +#
                 ((ch1 -# 0x80#) `uncheckedIShiftL#` 6#)  +#
                  (ch2 -# 0x80#)),
           3# #)

     | isTrue# ((ch0 >=# 0xF0#) `andI#` (ch0 <=# 0xF8#)) ->
        let !ch1 = word2Int# (indexWord8OffAddr# a# 1#) in
        if isTrue# ((ch1 <# 0x80#) `orI#` (ch1 >=# 0xC0#)) then fail 1# else
        let !ch2 = word2Int# (indexWord8OffAddr# a# 2#) in
        if isTrue# ((ch2 <# 0x80#) `orI#` (ch2 >=# 0xC0#)) then fail 2# else
        let !ch3 = word2Int# (indexWord8OffAddr# a# 3#) in
        if isTrue# ((ch3 <# 0x80#) `orI#` (ch3 >=# 0xC0#)) then fail 3# else
        (# chr# (((ch0 -# 0xF0#) `uncheckedIShiftL#` 18#) +#
                 ((ch1 -# 0x80#) `uncheckedIShiftL#` 12#) +#
                 ((ch2 -# 0x80#) `uncheckedIShiftL#` 6#)  +#
                  (ch3 -# 0x80#)),
           4# #)

      | otherwise -> fail 1#
  where
        -- all invalid sequences end up here:
        fail :: Int# -> (# Char#, Int# #)
        fail nBytes# = (# '\0'#, nBytes# #)
        -- '\xFFFD' would be the usual replacement character, but
        -- that's a valid symbol in Haskell, so will result in a
        -- confusing parse error later on.  Instead we use '\0' which
        -- will signal a lexer error immediately.

utf8DecodeChar :: Ptr Word8 -> (Char, Int)
utf8DecodeChar (Ptr a#) =
  case utf8DecodeChar# a# of (# c#, nBytes# #) -> ( C# c#, I# nBytes# )

-- UTF-8 is cleverly designed so that we can always figure out where
-- the start of the current character is, given any position in a
-- stream.  This function finds the start of the previous character,
-- assuming there *is* a previous character.
utf8PrevChar :: Ptr Word8 -> IO (Ptr Word8)
utf8PrevChar p = utf8CharStart (p `plusPtr` (-1))

utf8CharStart :: Ptr Word8 -> IO (Ptr Word8)
utf8CharStart p = go p
 where go p = do w <- peek p
                 if w >= 0x80 && w < 0xC0
                        then go (p `plusPtr` (-1))
                        else return p

utf8DecodeByteString :: ByteString -> [Char]
utf8DecodeByteString (BS.PS ptr offset len)
  = utf8DecodeStringLazy ptr offset len

utf8DecodeStringLazy :: ForeignPtr Word8 -> Int -> Int -> [Char]
utf8DecodeStringLazy fptr offset len
  = unsafeDupablePerformIO $ unpack start
  where
    !start = unsafeForeignPtrToPtr fptr `plusPtr` offset
    !end = start `plusPtr` len

    unpack p
        | p >= end  = touchForeignPtr fptr >> return []
        | otherwise =
            case utf8DecodeChar# (unPtr p) of
                (# c#, nBytes# #) -> do
                  rest <- unsafeDupableInterleaveIO $ unpack (p `plusPtr#` nBytes#)
                  return (C# c# : rest)

countUTF8Chars :: Ptr Word8 -> Int -> IO Int
countUTF8Chars ptr len = go ptr 0
  where
        !end = ptr `plusPtr` len

        go p !n
           | p >= end = return n
           | otherwise  = do
                case utf8DecodeChar# (unPtr p) of
                  (# _, nBytes# #) -> go (p `plusPtr#` nBytes#) (n+1)

unPtr :: Ptr a -> Addr#
unPtr (Ptr a) = a

plusPtr# :: Ptr a -> Int# -> Ptr a
plusPtr# ptr nBytes# = ptr `plusPtr` (I# nBytes#)

utf8EncodeChar :: Char -> Ptr Word8 -> IO (Ptr Word8)
utf8EncodeChar c ptr =
  let x = ord c in
  case () of
    _ | x > 0 && x <= 0x007f -> do
          poke ptr (fromIntegral x)
          return (ptr `plusPtr` 1)
        -- NB. '\0' is encoded as '\xC0\x80', not '\0'.  This is so that we
        -- can have 0-terminated UTF-8 strings (see GHC.Base.unpackCStringUtf8).
      | x <= 0x07ff -> do
          poke ptr (fromIntegral (0xC0 .|. ((x `shiftR` 6) .&. 0x1F)))
          pokeElemOff ptr 1 (fromIntegral (0x80 .|. (x .&. 0x3F)))
          return (ptr `plusPtr` 2)
      | x <= 0xffff -> do
          poke ptr (fromIntegral (0xE0 .|. (x `shiftR` 12) .&. 0x0F))
          pokeElemOff ptr 1 (fromIntegral (0x80 .|. (x `shiftR` 6) .&. 0x3F))
          pokeElemOff ptr 2 (fromIntegral (0x80 .|. (x .&. 0x3F)))
          return (ptr `plusPtr` 3)
      | otherwise -> do
          poke ptr (fromIntegral (0xF0 .|. (x `shiftR` 18)))
          pokeElemOff ptr 1 (fromIntegral (0x80 .|. ((x `shiftR` 12) .&. 0x3F)))
          pokeElemOff ptr 2 (fromIntegral (0x80 .|. ((x `shiftR` 6) .&. 0x3F)))
          pokeElemOff ptr 3 (fromIntegral (0x80 .|. (x .&. 0x3F)))
          return (ptr `plusPtr` 4)

utf8EncodeString :: Ptr Word8 -> String -> IO ()
utf8EncodeString ptr str = go ptr str
  where go !_   []     = return ()
        go ptr (c:cs) = do
          ptr' <- utf8EncodeChar c ptr
          go ptr' cs

utf8EncodedLength :: String -> Int
utf8EncodedLength str = go 0 str
  where go !n [] = n
        go n (c:cs)
          | ord c > 0 && ord c <= 0x007f = go (n+1) cs
          | ord c <= 0x07ff = go (n+2) cs
          | ord c <= 0xffff = go (n+3) cs
          | otherwise       = go (n+4) cs

-- -----------------------------------------------------------------------------
-- The Z-encoding

{-
This is the main name-encoding and decoding function.  It encodes any
string into a string that is acceptable as a C name.  This is done
right before we emit a symbol name into the compiled C or asm code.
Z-encoding of strings is cached in the FastString interface, so we
never encode the same string more than once.

The basic encoding scheme is this.

* Tuples (,,,) are coded as Z3T

* Alphabetic characters (upper and lower) and digits
        all translate to themselves;
        except 'Z', which translates to 'ZZ'
        and    'z', which translates to 'zz'
  We need both so that we can preserve the variable/tycon distinction

* Most other printable characters translate to 'zx' or 'Zx' for some
        alphabetic character x

* The others translate as 'znnnU' where 'nnn' is the decimal number
        of the character

        Before          After
        --------------------------
        Trak            Trak
        foo_wib         foozuwib
        >               zg
        >1              zg1
        foo#            foozh
        foo##           foozhzh
        foo##1          foozhzh1
        fooZ            fooZZ
        :+              ZCzp
        ()              Z0T     0-tuple
        (,,,,)          Z5T     5-tuple
        (# #)           Z1H     unboxed 1-tuple (note the space)
        (#,,,,#)        Z5H     unboxed 5-tuple
                (NB: There is no Z1T nor Z0H.)
-}

type UserString = String        -- As the user typed it
type EncodedString = String     -- Encoded form

zEncodeString :: UserString -> EncodedString
zEncodeString cs = case maybe_tuple cs of
                Just n  -> n            -- Tuples go to Z2T etc
                Nothing -> go cs
          where
                go []     = []
                go (c:cs) = encode_digit_ch c ++ go' cs
                go' []     = []
                go' (c:cs) = encode_ch c ++ go' cs

unencodedChar :: Char -> Bool   -- True for chars that don't need encoding
unencodedChar 'Z' = False
unencodedChar 'z' = False
unencodedChar c   =  c >= 'a' && c <= 'z'
                  || c >= 'A' && c <= 'Z'
                  || c >= '0' && c <= '9'

-- If a digit is at the start of a symbol then we need to encode it.
-- Otherwise package names like 9pH-0.1 give linker errors.
encode_digit_ch :: Char -> EncodedString
encode_digit_ch c | c >= '0' && c <= '9' = encode_as_unicode_char c
encode_digit_ch c | otherwise            = encode_ch c

encode_ch :: Char -> EncodedString
encode_ch c | unencodedChar c = [c]     -- Common case first

-- Constructors
encode_ch '('  = "ZL"   -- Needed for things like (,), and (->)
encode_ch ')'  = "ZR"   -- For symmetry with (
encode_ch '['  = "ZM"
encode_ch ']'  = "ZN"
encode_ch ':'  = "ZC"
encode_ch 'Z'  = "ZZ"

-- Variables
encode_ch 'z'  = "zz"
encode_ch '&'  = "za"
encode_ch '|'  = "zb"
encode_ch '^'  = "zc"
encode_ch '$'  = "zd"
encode_ch '='  = "ze"
encode_ch '>'  = "zg"
encode_ch '#'  = "zh"
encode_ch '.'  = "zi"
encode_ch '<'  = "zl"
encode_ch '-'  = "zm"
encode_ch '!'  = "zn"
encode_ch '+'  = "zp"
encode_ch '\'' = "zq"
encode_ch '\\' = "zr"
encode_ch '/'  = "zs"
encode_ch '*'  = "zt"
encode_ch '_'  = "zu"
encode_ch '%'  = "zv"
encode_ch c    = encode_as_unicode_char c

encode_as_unicode_char :: Char -> EncodedString
encode_as_unicode_char c = 'z' : if isDigit (head hex_str) then hex_str
                                                           else '0':hex_str
  where hex_str = showHex (ord c) "U"
  -- ToDo: we could improve the encoding here in various ways.
  -- eg. strings of unicode characters come out as 'z1234Uz5678U', we
  -- could remove the 'U' in the middle (the 'z' works as a separator).

zDecodeString :: EncodedString -> UserString
zDecodeString [] = []
zDecodeString ('Z' : d : rest)
  | isDigit d = decode_tuple   d rest
  | otherwise = decode_upper   d : zDecodeString rest
zDecodeString ('z' : d : rest)
  | isDigit d = decode_num_esc d rest
  | otherwise = decode_lower   d : zDecodeString rest
zDecodeString (c   : rest) = c : zDecodeString rest

decode_upper, decode_lower :: Char -> Char

decode_upper 'L' = '('
decode_upper 'R' = ')'
decode_upper 'M' = '['
decode_upper 'N' = ']'
decode_upper 'C' = ':'
decode_upper 'Z' = 'Z'
decode_upper ch  = {-pprTrace "decode_upper" (char ch)-} ch

decode_lower 'z' = 'z'
decode_lower 'a' = '&'
decode_lower 'b' = '|'
decode_lower 'c' = '^'
decode_lower 'd' = '$'
decode_lower 'e' = '='
decode_lower 'g' = '>'
decode_lower 'h' = '#'
decode_lower 'i' = '.'
decode_lower 'l' = '<'
decode_lower 'm' = '-'
decode_lower 'n' = '!'
decode_lower 'p' = '+'
decode_lower 'q' = '\''
decode_lower 'r' = '\\'
decode_lower 's' = '/'
decode_lower 't' = '*'
decode_lower 'u' = '_'
decode_lower 'v' = '%'
decode_lower ch  = {-pprTrace "decode_lower" (char ch)-} ch

-- Characters not having a specific code are coded as z224U (in hex)
decode_num_esc :: Char -> EncodedString -> UserString
decode_num_esc d rest
  = go (digitToInt d) rest
  where
    go n (c : rest) | isHexDigit c = go (16*n + digitToInt c) rest
    go n ('U' : rest)           = chr n : zDecodeString rest
    go n other = error ("decode_num_esc: " ++ show n ++  ' ':other)

decode_tuple :: Char -> EncodedString -> UserString
decode_tuple d rest
  = go (digitToInt d) rest
  where
        -- NB. recurse back to zDecodeString after decoding the tuple, because
        -- the tuple might be embedded in a longer name.
    go n (c : rest) | isDigit c = go (10*n + digitToInt c) rest
    go 0 ('T':rest)     = "()" ++ zDecodeString rest
    go n ('T':rest)     = '(' : replicate (n-1) ',' ++ ")" ++ zDecodeString rest
    go 1 ('H':rest)     = "(# #)" ++ zDecodeString rest
    go n ('H':rest)     = '(' : '#' : replicate (n-1) ',' ++ "#)" ++ zDecodeString rest
    go n other = error ("decode_tuple: " ++ show n ++ ' ':other)

{-
Tuples are encoded as
        Z3T or Z3H
for 3-tuples or unboxed 3-tuples respectively.  No other encoding starts
        Z<digit>

* "(# #)" is the tycon for an unboxed 1-tuple (not 0-tuple)
  There are no unboxed 0-tuples.

* "()" is the tycon for a boxed 0-tuple.
  There are no boxed 1-tuples.
-}

maybe_tuple :: UserString -> Maybe EncodedString

maybe_tuple "(# #)" = Just("Z1H")
maybe_tuple ('(' : '#' : cs) = case count_commas (0::Int) cs of
                                 (n, '#' : ')' : _) -> Just ('Z' : shows (n+1) "H")
                                 _                  -> Nothing
maybe_tuple "()" = Just("Z0T")
maybe_tuple ('(' : cs)       = case count_commas (0::Int) cs of
                                 (n, ')' : _) -> Just ('Z' : shows (n+1) "T")
                                 _            -> Nothing
maybe_tuple _                = Nothing

count_commas :: Int -> String -> (Int, String)
count_commas n (',' : cs) = count_commas (n+1) cs
count_commas n cs         = (n,cs)

{-
************************************************************************
*                                                                      *
                        Base 62
*                                                                      *
************************************************************************

Note [Base 62 encoding 128-bit integers]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Instead of base-62 encoding a single 128-bit integer
(ceil(21.49) characters), we'll base-62 a pair of 64-bit integers
(2 * ceil(10.75) characters).  Luckily for us, it's the same number of
characters!
-}

--------------------------------------------------------------------------
-- Base 62

-- The base-62 code is based off of 'locators'
-- ((c) Operational Dynamics Consulting, BSD3 licensed)

-- | Size of a 64-bit word when written as a base-62 string
word64Base62Len :: Int
word64Base62Len = 11

-- | Converts a 64-bit word into a base-62 string
toBase62Padded :: Word64 -> String
toBase62Padded w = pad ++ str
  where
    pad = replicate len '0'
    len = word64Base62Len - length str -- 11 == ceil(64 / lg 62)
    str = toBase62 w

toBase62 :: Word64 -> String
toBase62 w = showIntAtBase 62 represent w ""
  where
    represent :: Int -> Char
    represent x
        | x < 10 = Char.chr (48 + x)
        | x < 36 = Char.chr (65 + x - 10)
        | x < 62 = Char.chr (97 + x - 36)
        | otherwise = error "represent (base 62): impossible!"


-- | A tiny wrapper around 'IntSet.IntSet' for representing sets of 'Enum'
-- things.
module EnumSet
    ( EnumSet
    , member
    , insert
    , delete
    , toList
    , fromList
    , empty
    ) where

import GhcPrelude

import qualified Data.IntSet as IntSet

newtype EnumSet a = EnumSet IntSet.IntSet

member :: Enum a => a -> EnumSet a -> Bool
member x (EnumSet s) = IntSet.member (fromEnum x) s

insert :: Enum a => a -> EnumSet a -> EnumSet a
insert x (EnumSet s) = EnumSet $ IntSet.insert (fromEnum x) s

delete :: Enum a => a -> EnumSet a -> EnumSet a
delete x (EnumSet s) = EnumSet $ IntSet.delete (fromEnum x) s

toList :: Enum a => EnumSet a -> [a]
toList (EnumSet s) = map toEnum $ IntSet.toList s

fromList :: Enum a => [a] -> EnumSet a
fromList = EnumSet . IntSet.fromList . map fromEnum

empty :: EnumSet a
empty = EnumSet IntSet.empty


{-# OPTIONS_GHC -fno-warn-deprecations #-}
module Exception
    (
    module Control.Exception,
    module Exception
    )
    where

import GhcPrelude

import Control.Exception
import Control.Monad.IO.Class

catchIO :: IO a -> (IOException -> IO a) -> IO a
catchIO = Control.Exception.catch

handleIO :: (IOException -> IO a) -> IO a -> IO a
handleIO = flip catchIO

tryIO :: IO a -> IO (Either IOException a)
tryIO = try

-- | A monad that can catch exceptions.  A minimal definition
-- requires a definition of 'gcatch'.
--
-- Implementations on top of 'IO' should implement 'gmask' to
-- eventually call the primitive 'Control.Exception.mask'.
-- These are used for
-- implementations that support asynchronous exceptions.  The default
-- implementations of 'gbracket' and 'gfinally' use 'gmask'
-- thus rarely require overriding.
--
class MonadIO m => ExceptionMonad m where

  -- | Generalised version of 'Control.Exception.catch', allowing an arbitrary
  -- exception handling monad instead of just 'IO'.
  gcatch :: Exception e => m a -> (e -> m a) -> m a

  -- | Generalised version of 'Control.Exception.mask_', allowing an arbitrary
  -- exception handling monad instead of just 'IO'.
  gmask :: ((m a -> m a) -> m b) -> m b

  -- | Generalised version of 'Control.Exception.bracket', allowing an arbitrary
  -- exception handling monad instead of just 'IO'.
  gbracket :: m a -> (a -> m b) -> (a -> m c) -> m c

  -- | Generalised version of 'Control.Exception.finally', allowing an arbitrary
  -- exception handling monad instead of just 'IO'.
  gfinally :: m a -> m b -> m a

  gbracket before after thing =
    gmask $ \restore -> do
      a <- before
      r <- restore (thing a) `gonException` after a
      _ <- after a
      return r

  a `gfinally` sequel =
    gmask $ \restore -> do
      r <- restore a `gonException` sequel
      _ <- sequel
      return r

instance ExceptionMonad IO where
  gcatch    = Control.Exception.catch
  gmask f   = mask (\x -> f x)

gtry :: (ExceptionMonad m, Exception e) => m a -> m (Either e a)
gtry act = gcatch (act >>= \a -> return (Right a))
                  (\e -> return (Left e))

-- | Generalised version of 'Control.Exception.handle', allowing an arbitrary
-- exception handling monad instead of just 'IO'.
ghandle :: (ExceptionMonad m, Exception e) => (e -> m a) -> m a -> m a
ghandle = flip gcatch

-- | Always executes the first argument.  If this throws an exception the
-- second argument is executed and the exception is raised again.
gonException :: (ExceptionMonad m) => m a -> m b -> m a
gonException ioA cleanup = ioA `gcatch` \e ->
                             do _ <- cleanup
                                liftIO $ throwIO (e :: SomeException)


{-
(c) The University of Glasgow, 2000-2006
-}

{-# LANGUAGE CPP, MagicHash, UnboxedTuples #-}

module FastFunctions (
    inlinePerformIO,
  ) where

#include "HsVersions.h"

import GhcPrelude ()

import GHC.Exts
import GHC.IO   (IO(..))

-- Just like unsafeDupablePerformIO, but we inline it.
{-# INLINE inlinePerformIO #-}
inlinePerformIO :: IO a -> a
inlinePerformIO (IO m) = case m realWorld# of (# _, r #)   -> r


{-# LANGUAGE BangPatterns, MagicHash, UnboxedTuples #-}
{-# OPTIONS_GHC -O2 #-}
-- We always optimise this, otherwise performance of a non-optimised
-- compiler is severely affected
--
-- (c) The University of Glasgow 2002-2006
--
-- Unboxed mutable Ints

module FastMutInt(
        FastMutInt, newFastMutInt,
        readFastMutInt, writeFastMutInt,

        FastMutPtr, newFastMutPtr,
        readFastMutPtr, writeFastMutPtr
  ) where

import GhcPrelude

import Data.Bits
import GHC.Base
import GHC.Ptr

newFastMutInt :: IO FastMutInt
readFastMutInt :: FastMutInt -> IO Int
writeFastMutInt :: FastMutInt -> Int -> IO ()

newFastMutPtr :: IO FastMutPtr
readFastMutPtr :: FastMutPtr -> IO (Ptr a)
writeFastMutPtr :: FastMutPtr -> Ptr a -> IO ()

data FastMutInt = FastMutInt (MutableByteArray# RealWorld)

newFastMutInt = IO $ \s ->
  case newByteArray# size s of { (# s, arr #) ->
  (# s, FastMutInt arr #) }
  where !(I# size) = finiteBitSize (0 :: Int)

readFastMutInt (FastMutInt arr) = IO $ \s ->
  case readIntArray# arr 0# s of { (# s, i #) ->
  (# s, I# i #) }

writeFastMutInt (FastMutInt arr) (I# i) = IO $ \s ->
  case writeIntArray# arr 0# i s of { s ->
  (# s, () #) }

data FastMutPtr = FastMutPtr (MutableByteArray# RealWorld)

newFastMutPtr = IO $ \s ->
  case newByteArray# size s of { (# s, arr #) ->
  (# s, FastMutPtr arr #) }
  -- GHC assumes 'sizeof (Int) == sizeof (Ptr a)'
  where !(I# size) = finiteBitSize (0 :: Int)

readFastMutPtr (FastMutPtr arr) = IO $ \s ->
  case readAddrArray# arr 0# s of { (# s, i #) ->
  (# s, Ptr i #) }

writeFastMutPtr (FastMutPtr arr) (Ptr i) = IO $ \s ->
  case writeAddrArray# arr 0# i s of { s ->
  (# s, () #) }


-- (c) The University of Glasgow, 1997-2006

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

-- |
-- There are two principal string types used internally by GHC:
--
-- ['FastString']
--
--   * 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'.
--
-- ['PtrString']
--
--   * Pointer and size of a Latin-1 encoded string.
--   * Practically no operations.
--   * Outputing them is fast.
--   * Generated by 'sLit'.
--   * Turn into 'Outputable.SDoc' with 'Outputable.ptext'
--   * 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.
--
-- Use 'PtrString' unless you want the facilities of 'FastString'.
module FastString
       (
        -- * ByteString
        bytesFS,            -- :: FastString -> ByteString
        fastStringToByteString, -- = bytesFS (kept for haddock)
        mkFastStringByteString,
        fastZStringToByteString,
        unsafeMkByteString,

        -- * FastZString
        FastZString,
        hPutFZS,
        zString,
        lengthFZS,

        -- * FastStrings
        FastString(..),     -- not abstract, for now.

        -- ** Construction
        fsLit,
        mkFastString,
        mkFastStringBytes,
        mkFastStringByteList,
        mkFastStringForeignPtr,
        mkFastString#,

        -- ** Deconstruction
        unpackFS,           -- :: FastString -> String

        -- ** Encoding
        zEncodeFS,

        -- ** Operations
        uniqueOfFS,
        lengthFS,
        nullFS,
        appendFS,
        headFS,
        tailFS,
        concatFS,
        consFS,
        nilFS,
        isUnderscoreFS,

        -- ** Outputing
        hPutFS,

        -- ** Internal
        getFastStringTable,
        hasZEncoding,

        -- * PtrStrings
        PtrString (..),

        -- ** Construction
        sLit,
        mkPtrString#,
        mkPtrString,

        -- ** Deconstruction
        unpackPtrString,

        -- ** Operations
        lengthPS
       ) where

#include "HsVersions.h"

import GhcPrelude as Prelude

import Encoding
import FastFunctions
import Panic
import Util

import Control.Concurrent.MVar
import Control.DeepSeq
import Control.Monad
import Data.ByteString (ByteString)
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
import Foreign.C
import GHC.Exts
import System.IO
import Data.Data
import Data.IORef
import Data.Maybe       ( isJust )
import Data.Char
import Data.Semigroup as Semi

import GHC.IO

import Foreign

#if STAGE >= 2
import GHC.Conc.Sync    (sharedCAF)
#endif

import GHC.Base         ( unpackCString#, unpackNBytes# )

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

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

fastZStringToByteString :: FastZString -> ByteString
fastZStringToByteString (FastZString bs) = bs

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

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

-- -----------------------------------------------------------------------------

newtype FastZString = FastZString ByteString
  deriving NFData

hPutFZS :: Handle -> FastZString -> IO ()
hPutFZS handle (FastZString bs) = BS.hPut handle bs

zString :: FastZString -> String
zString (FastZString bs) =
    inlinePerformIO $ BS.unsafeUseAsCStringLen bs peekCAStringLen

lengthFZS :: FastZString -> Int
lengthFZS (FastZString bs) = BS.length bs

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

-- -----------------------------------------------------------------------------

{-|
A 'FastString' is an array of bytes, hashed to support fast O(1)
comparison.  It is also associated with a character encoding, so that
we know how to convert a 'FastString' to the local encoding, or to the
Z-encoding used by the compiler internally.

'FastString's support a memoized conversion to the Z-encoding via zEncodeFS.
-}

data FastString = FastString {
      uniq    :: {-# UNPACK #-} !Int, -- unique id
      n_chars :: {-# UNPACK #-} !Int, -- number of chars
      fs_bs   :: {-# UNPACK #-} !ByteString,
      fs_ref  :: {-# UNPACK #-} !(IORef (Maybe FastZString))
  }

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

instance Ord FastString where
    -- Compares lexicographically, not by unique
    a <= b = case cmpFS a b of { LT -> True;  EQ -> True;  GT -> False }
    a <  b = case cmpFS a b of { LT -> True;  EQ -> False; GT -> False }
    a >= b = case cmpFS a b of { LT -> False; EQ -> True;  GT -> True  }
    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
    compare a b = cmpFS a b

instance IsString FastString where
    fromString = fsLit

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

instance Monoid FastString where
    mempty = nilFS
    mappend = (Semi.<>)
    mconcat = concatFS

instance Show FastString where
   show fs = show (unpackFS fs)

instance Data FastString where
  -- don't traverse?
  toConstr _   = abstractConstr "FastString"
  gunfold _ _  = error "gunfold"
  dataTypeOf _ = mkNoRepType "FastString"

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

foreign import ccall unsafe "memcmp"
  memcmp :: Ptr a -> Ptr b -> Int -> IO Int

-- -----------------------------------------------------------------------------
-- Construction

{-
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.

See Note [Updating the FastString table] on how it's updated.
-}
data FastStringTable = FastStringTable
  {-# UNPACK #-} !(IORef Int) -- the unique ID counter shared with all buckets
  (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#
  uid <- newIORef 603979776 -- ord '$' * 0x01000000
  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
          (# s4#, segments# #) -> (# s4#, FastStringTable uid segments# #)

  -- use the support wired into the RTS to share this CAF among all images of
  -- libHSghc
#if STAGE < 2
  return tab
#else
  sharedCAF tab getOrSetLibHSghcFastStringTable

-- from the RTS; thus we cannot use this mechanism when STAGE<2; the previous
-- 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.

-}

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

{- Note [Updating the FastString table]

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.

The procedure goes like this:

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:

   * Resize the segment by doubling the number of buckets when the number of
     FastStrings in this segment grows beyond the threshold.
   * 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.
-}

mkFastStringWith :: (Int -> IO FastString) -> Ptr Word8 -> Int -> IO FastString
mkFastStringWith mk_fs !ptr !len = do
  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
      -- 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
      n <- get_uid
      new_fs <- mk_fs n
      withMVar lock $ \_ -> insert new_fs
  where
    !(FastStringTable uid segments#) = stringTable
    get_uid = atomicModifyIORef' uid $ \n -> (n+1,n)

    !(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

mkFastStringBytes :: Ptr Word8 -> Int -> FastString
mkFastStringBytes !ptr !len =
    -- NB: Might as well use unsafeDupablePerformIO, since mkFastStringWith is
    -- idempotent.
    unsafeDupablePerformIO $
        mkFastStringWith (copyNewFastString ptr len) ptr len

-- | 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
mkFastStringForeignPtr ptr !fp len
    = mkFastStringWith (mkNewFastString fp ptr len) ptr len

-- | 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.
mkFastStringByteString :: ByteString -> FastString
mkFastStringByteString bs =
    inlinePerformIO $
      BS.unsafeUseAsCStringLen bs $ \(ptr, len) -> do
        let ptr' = castPtr ptr
        mkFastStringWith (mkNewFastStringByteString bs ptr' len) ptr' len

-- | Creates a UTF-8 encoded 'FastString' from a 'String'
mkFastString :: String -> FastString
mkFastString str =
  inlinePerformIO $ do
    let l = utf8EncodedLength str
    buf <- mallocForeignPtrBytes l
    withForeignPtr buf $ \ptr -> do
      utf8EncodeString ptr str
      mkFastStringForeignPtr ptr buf l

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

-- | Creates a Z-encoded 'FastString' from a 'String'
mkZFastString :: String -> FastZString
mkZFastString = mkFastZStringString

mkNewFastString :: ForeignPtr Word8 -> Ptr Word8 -> Int -> Int
                -> IO FastString
mkNewFastString fp ptr len uid = do
  ref <- newIORef Nothing
  n_chars <- countUTF8Chars ptr len
  return (FastString uid n_chars (BS.fromForeignPtr fp 0 len) ref)

mkNewFastStringByteString :: ByteString -> Ptr Word8 -> Int -> Int
                          -> IO FastString
mkNewFastStringByteString bs ptr len uid = do
  ref <- newIORef Nothing
  n_chars <- countUTF8Chars ptr len
  return (FastString uid n_chars bs ref)

copyNewFastString :: Ptr Word8 -> Int -> Int -> IO FastString
copyNewFastString ptr len uid = do
  fp <- copyBytesToForeignPtr ptr len
  ref <- newIORef Nothing
  n_chars <- countUTF8Chars ptr len
  return (FastString uid n_chars (BS.fromForeignPtr fp 0 len) ref)

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

cmpStringPrefix :: Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
cmpStringPrefix ptr1 ptr2 len =
 do r <- memcmp ptr1 ptr2 len
    return (r == 0)

hashStr  :: Ptr Word8 -> Int -> Int
 -- use the Addr to produce a hash value between 0 & m (inclusive)
hashStr (Ptr a#) (I# len#) = loop 0# 0#
   where
    loop h n | isTrue# (n ==# len#) = I# h
             | otherwise  = loop h2 (n +# 1#)
          where
            !c = ord# (indexCharOffAddr# a# n)
            !h2 = (h *# 16777619#) `xorI#` c

-- -----------------------------------------------------------------------------
-- Operations

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

-- | Returns @True@ if this 'FastString' is not Z-encoded but already has
-- a Z-encoding cached (used in producing stats).
hasZEncoding :: FastString -> Bool
hasZEncoding (FastString _ _ _ ref) =
      inlinePerformIO $ do
        m <- readIORef ref
        return (isJust m)

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

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

-- | Returns a Z-encoded version of a 'FastString'.  This might be the
-- original, if it was already Z-encoded.  The first time this
-- function is applied to a particular 'FastString', the results are
-- memoized.
--
zEncodeFS :: FastString -> FastZString
zEncodeFS fs@(FastString _ _ _ ref) =
      inlinePerformIO $ do
        m <- readIORef ref
        case m of
          Just zfs -> return zfs
          Nothing -> do
            atomicModifyIORef' ref $ \m' -> case m' of
              Nothing  -> let zfs = mkZFastString (zEncodeString (unpackFS fs))
                          in (Just zfs, zfs)
              Just zfs -> (m', zfs)

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

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

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

tailFS :: FastString -> FastString
tailFS (FastString _ 0 _ _) = panic "tailFS: Empty FastString"
tailFS (FastString _ _ bs _) =
    inlinePerformIO $ BS.unsafeUseAsCString bs $ \ptr ->
    do let (_, n) = utf8DecodeChar (castPtr ptr)
       return $! mkFastStringByteString (BS.drop n bs)

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

uniqueOfFS :: FastString -> Int
uniqueOfFS (FastString u _ _ _) = u

nilFS :: FastString
nilFS = mkFastString ""

isUnderscoreFS :: FastString -> Bool
isUnderscoreFS fs = fs == fsLit "_"

-- -----------------------------------------------------------------------------
-- Stats

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
    !(FastStringTable _ segments#) = stringTable

-- -----------------------------------------------------------------------------
-- Outputting 'FastString's

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

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

-- -----------------------------------------------------------------------------
-- PtrStrings, here for convenience only.

-- | A 'PtrString' is a pointer to some array of Latin-1 encoded chars.
data PtrString = PtrString !(Ptr Word8) !Int

-- | Wrap an unboxed address into a 'PtrString'.
mkPtrString# :: Addr# -> PtrString
mkPtrString# a# = PtrString (Ptr a#) (ptrStrLength (Ptr a#))

-- | Encode a 'String' into a newly allocated 'PtrString' using Latin-1
-- encoding.  The original string must not contain non-Latin-1 characters
-- (above codepoint @0xff@).
{-# INLINE mkPtrString #-}
mkPtrString :: String -> PtrString
mkPtrString s =
 -- we don't use `unsafeDupablePerformIO` here to avoid potential memory leaks
 -- and because someone might be using `eqAddr#` to check for string equality.
 unsafePerformIO (do
   let len = length s
   p <- mallocBytes len
   let
     loop :: Int -> String -> IO ()
     loop !_ []    = return ()
     loop n (c:cs) = do
        pokeByteOff p n (fromIntegral (ord c) :: Word8)
        loop (1+n) cs
   loop 0 s
   return (PtrString p len)
 )

-- | 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#

-- | Return the length of a 'PtrString'
lengthPS :: PtrString -> Int
lengthPS (PtrString _ n) = n

-- -----------------------------------------------------------------------------
-- under the carpet

foreign import ccall unsafe "strlen"
  ptrStrLength :: Ptr Word8 -> Int

{-# NOINLINE sLit #-}
sLit :: String -> PtrString
sLit x  = mkPtrString x

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

{-# RULES "slit"
    forall x . sLit  (unpackCString# x) = mkPtrString#  x #-}
{-# RULES "fslit"
    forall x . fsLit (unpackCString# x) = mkFastString# x #-}


{-
%
% (c) The University of Glasgow 2006
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section[FastStringEnv]{@FastStringEnv@: FastString environments}
-}

module FastStringEnv (
        -- * FastString environments (maps)
        FastStringEnv,

        -- ** Manipulating these environments
        mkFsEnv,
        emptyFsEnv, unitFsEnv,
        extendFsEnv_C, extendFsEnv_Acc, extendFsEnv,
        extendFsEnvList, extendFsEnvList_C,
        filterFsEnv,
        plusFsEnv, plusFsEnv_C, alterFsEnv,
        lookupFsEnv, lookupFsEnv_NF, delFromFsEnv, delListFromFsEnv,
        elemFsEnv, mapFsEnv,

        -- * Deterministic FastString environments (maps)
        DFastStringEnv,

        -- ** Manipulating these environments
        mkDFsEnv, emptyDFsEnv, dFsEnvElts, lookupDFsEnv
    ) where

import GhcPrelude

import UniqFM
import UniqDFM
import Maybes
import FastString

-- | A non-deterministic set of FastStrings.
-- See Note [Deterministic UniqFM] in UniqDFM for explanation why it's not
-- deterministic and why it matters. Use DFastStringEnv if the set eventually
-- gets converted into a list or folded over in a way where the order
-- changes the generated code.
type FastStringEnv a = UniqFM a  -- Domain is FastString

emptyFsEnv         :: FastStringEnv a
mkFsEnv            :: [(FastString,a)] -> FastStringEnv a
alterFsEnv         :: (Maybe a-> Maybe a) -> FastStringEnv a -> FastString -> FastStringEnv a
extendFsEnv_C      :: (a->a->a) -> FastStringEnv a -> FastString -> a -> FastStringEnv a
extendFsEnv_Acc    :: (a->b->b) -> (a->b) -> FastStringEnv b -> FastString -> a -> FastStringEnv b
extendFsEnv        :: FastStringEnv a -> FastString -> a -> FastStringEnv a
plusFsEnv          :: FastStringEnv a -> FastStringEnv a -> FastStringEnv a
plusFsEnv_C        :: (a->a->a) -> FastStringEnv a -> FastStringEnv a -> FastStringEnv a
extendFsEnvList    :: FastStringEnv a -> [(FastString,a)] -> FastStringEnv a
extendFsEnvList_C  :: (a->a->a) -> FastStringEnv a -> [(FastString,a)] -> FastStringEnv a
delFromFsEnv       :: FastStringEnv a -> FastString -> FastStringEnv a
delListFromFsEnv   :: FastStringEnv a -> [FastString] -> FastStringEnv a
elemFsEnv          :: FastString -> FastStringEnv a -> Bool
unitFsEnv          :: FastString -> a -> FastStringEnv a
lookupFsEnv        :: FastStringEnv a -> FastString -> Maybe a
lookupFsEnv_NF     :: FastStringEnv a -> FastString -> a
filterFsEnv        :: (elt -> Bool) -> FastStringEnv elt -> FastStringEnv elt
mapFsEnv           :: (elt1 -> elt2) -> FastStringEnv elt1 -> FastStringEnv elt2

emptyFsEnv                = emptyUFM
unitFsEnv x y             = unitUFM x y
extendFsEnv x y z         = addToUFM x y z
extendFsEnvList x l       = addListToUFM x l
lookupFsEnv x y           = lookupUFM x y
alterFsEnv                = alterUFM
mkFsEnv     l             = listToUFM l
elemFsEnv x y             = elemUFM x y
plusFsEnv x y             = plusUFM x y
plusFsEnv_C f x y         = plusUFM_C f x y
extendFsEnv_C f x y z     = addToUFM_C f x y z
mapFsEnv f x              = mapUFM f x
extendFsEnv_Acc x y z a b = addToUFM_Acc x y z a b
extendFsEnvList_C x y z   = addListToUFM_C x y z
delFromFsEnv x y          = delFromUFM x y
delListFromFsEnv x y      = delListFromUFM x y
filterFsEnv x y           = filterUFM x y

lookupFsEnv_NF env n = expectJust "lookupFsEnv_NF" (lookupFsEnv env n)

-- Deterministic FastStringEnv
-- See Note [Deterministic UniqFM] in UniqDFM for explanation why we need
-- DFastStringEnv.

type DFastStringEnv a = UniqDFM a  -- Domain is FastString

emptyDFsEnv :: DFastStringEnv a
emptyDFsEnv = emptyUDFM

dFsEnvElts :: DFastStringEnv a -> [a]
dFsEnvElts = eltsUDFM

mkDFsEnv :: [(FastString,a)] -> DFastStringEnv a
mkDFsEnv l = listToUDFM l

lookupDFsEnv :: DFastStringEnv a -> FastString -> Maybe a
lookupDFsEnv = lookupUDFM


-- Some extra functions to extend Data.Map

module FiniteMap (
        insertList,
        insertListWith,
        deleteList,
        foldRight, foldRightWithKey
    ) where

import GhcPrelude

import Data.Map (Map)
import qualified Data.Map as Map

insertList :: Ord key => [(key,elt)] -> Map key elt -> Map key elt
insertList xs m = foldl' (\m (k, v) -> Map.insert k v m) m xs

insertListWith :: Ord key
               => (elt -> elt -> elt)
               -> [(key,elt)]
               -> Map key elt
               -> Map key elt
insertListWith f xs m0 = foldl' (\m (k, v) -> Map.insertWith f k v m) m0 xs

deleteList :: Ord key => [key] -> Map key elt -> Map key elt
deleteList ks m = foldl' (flip Map.delete) m ks

foldRight        :: (elt -> a -> a) -> a -> Map key elt -> a
foldRight        = Map.foldr
foldRightWithKey :: (key -> elt -> a -> a) -> a -> Map key elt -> a
foldRightWithKey = Map.foldrWithKey


{-
(c) Bartosz Nitka, Facebook 2015

Utilities for efficiently and deterministically computing free variables.

-}

{-# LANGUAGE BangPatterns #-}

module FV (
        -- * Deterministic free vars computations
        FV, InterestingVarFun,

        -- * Running the computations
        fvVarListVarSet, fvVarList, fvVarSet, fvDVarSet,

        -- ** Manipulating those computations
        unitFV,
        emptyFV,
        mkFVs,
        unionFV,
        unionsFV,
        delFV,
        delFVs,
        filterFV,
        mapUnionFV,
    ) where

import GhcPrelude

import Var
import VarSet

-- | Predicate on possible free variables: returns @True@ iff the variable is
-- interesting
type InterestingVarFun = Var -> Bool

-- Note [Deterministic FV]
-- ~~~~~~~~~~~~~~~~~~~~~~~
-- When computing free variables, the order in which you get them affects
-- the results of floating and specialization. If you use UniqFM to collect
-- them and then turn that into a list, you get them in nondeterministic
-- order as described in Note [Deterministic UniqFM] in UniqDFM.

-- A naive algorithm for free variables relies on merging sets of variables.
-- Merging costs O(n+m) for UniqFM and for UniqDFM there's an additional log
-- factor. It's cheaper to incrementally add to a list and use a set to check
-- for duplicates.
type FV = InterestingVarFun
             -- Used for filtering sets as we build them
          -> VarSet
             -- Locally bound variables
          -> ([Var], VarSet)
             -- List to preserve ordering and set to check for membership,
             -- so that the list doesn't have duplicates
             -- For explanation of why using `VarSet` is not deterministic see
             -- Note [Deterministic UniqFM] in UniqDFM.
          -> ([Var], VarSet)

-- Note [FV naming conventions]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- To get the performance and determinism that FV provides, FV computations
-- need to built up from smaller FV computations and then evaluated with
-- one of `fvVarList`, `fvDVarSet`, `fvVarListVarSet`. That means the functions
-- returning FV need to be exported.
--
-- The conventions are:
--
-- a) non-deterministic functions:
--   * a function that returns VarSet
--       e.g. `tyVarsOfType`
-- b) deterministic functions:
--   * a worker that returns FV
--       e.g. `tyFVsOfType`
--   * a function that returns [Var]
--       e.g. `tyVarsOfTypeList`
--   * a function that returns DVarSet
--       e.g. `tyVarsOfTypeDSet`
--
-- Where tyVarsOfType, tyVarsOfTypeList, tyVarsOfTypeDSet are implemented
-- in terms of the worker evaluated with fvVarSet, fvVarList, fvDVarSet
-- respectively.

-- | Run a free variable computation, returning a list of distinct free
-- variables in deterministic order and a non-deterministic set containing
-- those variables.
fvVarListVarSet :: FV ->  ([Var], VarSet)
fvVarListVarSet fv = fv (const True) emptyVarSet ([], emptyVarSet)

-- | Run a free variable computation, returning a list of distinct free
-- variables in deterministic order.
fvVarList :: FV -> [Var]
fvVarList = fst . fvVarListVarSet

-- | Run a free variable computation, returning a deterministic set of free
-- variables. Note that this is just a wrapper around the version that
-- returns a deterministic list. If you need a list you should use
-- `fvVarList`.
fvDVarSet :: FV -> DVarSet
fvDVarSet = mkDVarSet . fst . fvVarListVarSet

-- | Run a free variable computation, returning a non-deterministic set of
-- free variables. Don't use if the set will be later converted to a list
-- and the order of that list will impact the generated code.
fvVarSet :: FV -> VarSet
fvVarSet = snd . fvVarListVarSet

-- Note [FV eta expansion]
-- ~~~~~~~~~~~~~~~~~~~~~~~
-- Let's consider an eta-reduced implementation of freeVarsOf using FV:
--
-- freeVarsOf (App a b) = freeVarsOf a `unionFV` freeVarsOf b
--
-- If GHC doesn't eta-expand it, after inlining unionFV we end up with
--
-- freeVarsOf = \x ->
--   case x of
--     App a b -> \fv_cand in_scope acc ->
--       freeVarsOf a fv_cand in_scope $! freeVarsOf b fv_cand in_scope $! acc
--
-- which has to create a thunk, resulting in more allocations.
--
-- On the other hand if it is eta-expanded:
--
-- freeVarsOf (App a b) fv_cand in_scope acc =
--   (freeVarsOf a `unionFV` freeVarsOf b) fv_cand in_scope acc
--
-- after inlining unionFV we have:
--
-- freeVarsOf = \x fv_cand in_scope acc ->
--   case x of
--     App a b ->
--       freeVarsOf a fv_cand in_scope $! freeVarsOf b fv_cand in_scope $! acc
--
-- which saves allocations.
--
-- GHC when presented with knowledge about all the call sites, correctly
-- eta-expands in this case. Unfortunately due to the fact that freeVarsOf gets
-- exported to be composed with other functions, GHC doesn't have that
-- information and has to be more conservative here.
--
-- Hence functions that get exported and return FV need to be manually
-- eta-expanded. See also #11146.

-- | Add a variable - when free, to the returned free variables.
-- Ignores duplicates and respects the filtering function.
unitFV :: Id -> FV
unitFV var fv_cand in_scope acc@(have, haveSet)
  | var `elemVarSet` in_scope = acc
  | var `elemVarSet` haveSet = acc
  | fv_cand var = (var:have, extendVarSet haveSet var)
  | otherwise = acc
{-# INLINE unitFV #-}

-- | Return no free variables.
emptyFV :: FV
emptyFV _ _ acc = acc
{-# INLINE emptyFV #-}

-- | Union two free variable computations.
unionFV :: FV -> FV -> FV
unionFV fv1 fv2 fv_cand in_scope acc =
  fv1 fv_cand in_scope $! fv2 fv_cand in_scope $! acc
{-# INLINE unionFV #-}

-- | Mark the variable as not free by putting it in scope.
delFV :: Var -> FV -> FV
delFV var fv fv_cand !in_scope acc =
  fv fv_cand (extendVarSet in_scope var) acc
{-# INLINE delFV #-}

-- | Mark many free variables as not free.
delFVs :: VarSet -> FV -> FV
delFVs vars fv fv_cand !in_scope acc =
  fv fv_cand (in_scope `unionVarSet` vars) acc
{-# INLINE delFVs #-}

-- | Filter a free variable computation.
filterFV :: InterestingVarFun -> FV -> FV
filterFV fv_cand2 fv fv_cand1 in_scope acc =
  fv (\v -> fv_cand1 v && fv_cand2 v) in_scope acc
{-# INLINE filterFV #-}

-- | Map a free variable computation over a list and union the results.
mapUnionFV :: (a -> FV) -> [a] -> FV
mapUnionFV _f [] _fv_cand _in_scope acc = acc
mapUnionFV f (a:as) fv_cand in_scope acc =
  mapUnionFV f as fv_cand in_scope $! f a fv_cand in_scope $! acc
{-# INLINABLE mapUnionFV #-}

-- | Union many free variable computations.
unionsFV :: [FV] -> FV
unionsFV fvs fv_cand in_scope acc = mapUnionFV id fvs fv_cand in_scope acc
{-# INLINE unionsFV #-}

-- | Add multiple variables - when free, to the returned free variables.
-- Ignores duplicates and respects the filtering function.
mkFVs :: [Var] -> FV
mkFVs vars fv_cand in_scope acc =
  mapUnionFV unitFV vars fv_cand in_scope acc
{-# INLINE mkFVs #-}


{-# LANGUAGE CPP #-}

-- | Custom GHC "Prelude"
--
-- This module serves as a replacement for the "Prelude" module
-- and abstracts over differences between the bootstrapping
-- GHC version, and may also provide a common default vocabulary.

-- Every module in GHC
--   * Is compiled with -XNoImplicitPrelude
--   * Explicitly imports GhcPrelude

module GhcPrelude (module X) where

-- We export the 'Semigroup' class but w/o the (<>) operator to avoid
-- clashing with the (Outputable.<>) operator which is heavily used
-- through GHC's code-base.

import Prelude as X hiding ((<>))
import Data.Foldable as X (foldl')

{-
Note [Why do we import Prelude here?]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The files ghc-boot-th.cabal, ghc-boot.cabal, ghci.cabal and
ghc-heap.cabal contain the directive default-extensions:
NoImplicitPrelude. There are two motivations for this:
  - Consistency with the compiler directory, which enables
    NoImplicitPrelude;
  - Allows loading the above dependent packages with ghc-in-ghci,
    giving a smoother development experience when adding new
    extensions.
-}


-- | Types for the general graph colorer.
module GraphBase (
        Triv,
        Graph (..),
        initGraph,
        graphMapModify,

        Node  (..),     newNode,
)

where

import GhcPrelude

import UniqSet
import UniqFM

-- | A fn to check if a node is trivially colorable
--      For graphs who's color classes are disjoint then a node is 'trivially colorable'
--      when it has less neighbors and exclusions than available colors for that node.
--
--      For graph's who's color classes overlap, ie some colors alias other colors, then
--      this can be a bit more tricky. There is a general way to calculate this, but
--      it's likely be too slow for use in the code. The coloring algorithm takes
--      a canned function which can be optimised by the user to be specific to the
--      specific graph being colored.
--
--      for details, see  "A Generalised Algorithm for Graph-Coloring Register Allocation"
--                              Smith, Ramsey, Holloway - PLDI 2004.
--
type Triv k cls color
        =  cls                  -- the class of the node we're trying to color.
        -> UniqSet k            -- the node's neighbors.
        -> UniqSet color        -- the node's exclusions.
        -> Bool

-- | The Interference graph.
--      There used to be more fields, but they were turfed out in a previous revision.
--      maybe we'll want more later..
--
data Graph k cls color
        = Graph {
        -- | All active nodes in the graph.
          graphMap              :: UniqFM (Node k cls color)  }

-- | An empty graph.
initGraph :: Graph k cls color
initGraph
        = Graph
        { graphMap              = emptyUFM }

-- | Modify the finite map holding the nodes in the graph.
graphMapModify
        :: (UniqFM (Node k cls color) -> UniqFM (Node k cls color))
        -> Graph k cls color -> Graph k cls color

graphMapModify f graph
        = graph { graphMap      = f (graphMap graph) }

-- | Graph nodes.
--      Represents a thing that can conflict with another thing.
--      For the register allocater the nodes represent registers.
--
data Node k cls color
        = Node {
        -- | A unique identifier for this node.
          nodeId                :: k

        -- | The class of this node,
        --      determines the set of colors that can be used.
        , nodeClass             :: cls

        -- | The color of this node, if any.
        , nodeColor             :: Maybe color

        -- | Neighbors which must be colored differently to this node.
        , nodeConflicts         :: UniqSet k

        -- | Colors that cannot be used by this node.
        , nodeExclusions        :: UniqSet color

        -- | Colors that this node would prefer to be, in decending order.
        , nodePreference        :: [color]

        -- | Neighbors that this node would like to be colored the same as.
        , nodeCoalesce          :: UniqSet k }

-- | An empty node.
newNode :: k -> cls -> Node k cls color
newNode k cls
        = Node
        { nodeId                = k
        , nodeClass             = cls
        , nodeColor             = Nothing
        , nodeConflicts         = emptyUniqSet
        , nodeExclusions        = emptyUniqSet
        , nodePreference        = []
        , nodeCoalesce          = emptyUniqSet }


-- | Graph Coloring.
--      This is a generic graph coloring library, abstracted over the type of
--      the node keys, nodes and colors.
--

module GraphColor (
        module GraphBase,
        module GraphOps,
        module GraphPpr,
        colorGraph
)

where

import GhcPrelude

import GraphBase
import GraphOps
import GraphPpr

import Unique
import UniqFM
import UniqSet
import Outputable

import Data.Maybe
import Data.List

-- | Try to color a graph with this set of colors.
--      Uses Chaitin's algorithm to color the graph.
--      The graph is scanned for nodes which are deamed 'trivially colorable'. These nodes
--      are pushed onto a stack and removed from the graph.
--      Once this process is complete the graph can be colored by removing nodes from
--      the stack (ie in reverse order) and assigning them colors different to their neighbors.
--
colorGraph
        :: ( Uniquable  k, Uniquable cls,  Uniquable  color
           , Eq cls, Ord k
           , Outputable k, Outputable cls, Outputable color)
        => Bool                         -- ^ whether to do iterative coalescing
        -> Int                          -- ^ how many times we've tried to color this graph so far.
        -> UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
        -> Triv   k cls color           -- ^ fn to decide whether a node is trivially colorable.
        -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
        -> Graph  k cls color           -- ^ the graph to color.

        -> ( Graph k cls color          -- the colored graph.
           , UniqSet k                  -- the set of nodes that we couldn't find a color for.
           , UniqFM  k )                -- map of regs (r1 -> r2) that were coalesced
                                        --       r1 should be replaced by r2 in the source

colorGraph iterative spinCount colors triv spill graph0
 = let
        -- If we're not doing iterative coalescing then do an aggressive coalescing first time
        --      around and then conservative coalescing for subsequent passes.
        --
        --      Aggressive coalescing is a quick way to get rid of many reg-reg moves. However, if
        --      there is a lot of register pressure and we do it on every round then it can make the
        --      graph less colorable and prevent the algorithm from converging in a sensible number
        --      of cycles.
        --
        (graph_coalesced, kksCoalesce1)
         = if iterative
                then (graph0, [])
                else if spinCount == 0
                        then coalesceGraph True  triv graph0
                        else coalesceGraph False triv graph0

        -- run the scanner to slurp out all the trivially colorable nodes
        --      (and do coalescing if iterative coalescing is enabled)
        (ksTriv, ksProblems, kksCoalesce2)
                = colorScan iterative triv spill graph_coalesced

        -- If iterative coalescing is enabled, the scanner will coalesce the graph as does its business.
        --      We need to apply all the coalescences found by the scanner to the original
        --      graph before doing assignColors.
        --
        --      Because we've got the whole, non-pruned graph here we turn on aggressive coalecing
        --      to force all the (conservative) coalescences found during scanning.
        --
        (graph_scan_coalesced, _)
                = mapAccumL (coalesceNodes True triv) graph_coalesced kksCoalesce2

        -- color the trivially colorable nodes
        --      during scanning, keys of triv nodes were added to the front of the list as they were found
        --      this colors them in the reverse order, as required by the algorithm.
        (graph_triv, ksNoTriv)
                = assignColors colors graph_scan_coalesced ksTriv

        -- try and color the problem nodes
        --      problem nodes are the ones that were left uncolored because they weren't triv.
        --      theres a change we can color them here anyway.
        (graph_prob, ksNoColor)
                = assignColors colors graph_triv ksProblems

        -- if the trivially colorable nodes didn't color then something is probably wrong
        --      with the provided triv function.
        --
   in   if not $ null ksNoTriv
         then   pprPanic "colorGraph: trivially colorable nodes didn't color!" -- empty
                        (  empty
                        $$ text "ksTriv    = " <> ppr ksTriv
                        $$ text "ksNoTriv  = " <> ppr ksNoTriv
                        $$ text "colors    = " <> ppr colors
                        $$ empty
                        $$ dotGraph (\_ -> text "white") triv graph_triv)

         else   ( graph_prob
                , mkUniqSet ksNoColor   -- the nodes that didn't color (spills)
                , if iterative
                        then (listToUFM kksCoalesce2)
                        else (listToUFM kksCoalesce1))

-- | Scan through the conflict graph separating out trivially colorable and
--      potentially uncolorable (problem) nodes.
--
--      Checking whether a node is trivially colorable or not is a reasonably expensive operation,
--      so after a triv node is found and removed from the graph it's no good to return to the 'start'
--      of the graph and recheck a bunch of nodes that will probably still be non-trivially colorable.
--
--      To ward against this, during each pass through the graph we collect up a list of triv nodes
--      that were found, and only remove them once we've finished the pass. The more nodes we can delete
--      at once the more likely it is that nodes we've already checked will become trivially colorable
--      for the next pass.
--
--      TODO:   add work lists to finding triv nodes is easier.
--              If we've just scanned the graph, and removed triv nodes, then the only
--              nodes that we need to rescan are the ones we've removed edges from.

colorScan
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool                         -- ^ whether to do iterative coalescing
        -> Triv k cls color             -- ^ fn to decide whether a node is trivially colorable
        -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
        -> Graph k cls color            -- ^ the graph to scan

        -> ([k], [k], [(k, k)])         --  triv colorable nodes, problem nodes, pairs of nodes to coalesce

colorScan iterative triv spill graph
        = colorScan_spin iterative triv spill graph [] [] []

colorScan_spin
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool
        -> Triv k cls color
        -> (Graph k cls color -> k)
        -> Graph k cls color
        -> [k]
        -> [k]
        -> [(k, k)]
        -> ([k], [k], [(k, k)])

colorScan_spin iterative triv spill graph
        ksTriv ksSpill kksCoalesce

        -- if the graph is empty then we're done
        | isNullUFM $ graphMap graph
        = (ksTriv, ksSpill, reverse kksCoalesce)

        -- Simplify:
        --      Look for trivially colorable nodes.
        --      If we can find some then remove them from the graph and go back for more.
        --
        | nsTrivFound@(_:_)
                <-  scanGraph   (\node -> triv  (nodeClass node) (nodeConflicts node) (nodeExclusions node)

                                  -- for iterative coalescing we only want non-move related
                                  --    nodes here
                                  && (not iterative || isEmptyUniqSet (nodeCoalesce node)))
                        $ graph

        , ksTrivFound   <- map nodeId nsTrivFound
        , graph2        <- foldr (\k g -> let Just g' = delNode k g
                                          in  g')
                                graph ksTrivFound

        = colorScan_spin iterative triv spill graph2
                (ksTrivFound ++ ksTriv)
                ksSpill
                kksCoalesce

        -- Coalesce:
        --      If we're doing iterative coalescing and no triv nodes are available
        --      then it's time for a coalescing pass.
        | iterative
        = case coalesceGraph False triv graph of

                -- we were able to coalesce something
                --      go back to Simplify and see if this frees up more nodes to be trivially colorable.
                (graph2, kksCoalesceFound@(_:_))
                 -> colorScan_spin iterative triv spill graph2
                        ksTriv ksSpill (reverse kksCoalesceFound ++ kksCoalesce)

                -- Freeze:
                -- nothing could be coalesced (or was triv),
                --      time to choose a node to freeze and give up on ever coalescing it.
                (graph2, [])
                 -> case freezeOneInGraph graph2 of

                        -- we were able to freeze something
                        --      hopefully this will free up something for Simplify
                        (graph3, True)
                         -> colorScan_spin iterative triv spill graph3
                                ksTriv ksSpill kksCoalesce

                        -- we couldn't find something to freeze either
                        --      time for a spill
                        (graph3, False)
                         -> colorScan_spill iterative triv spill graph3
                                ksTriv ksSpill kksCoalesce

        -- spill time
        | otherwise
        = colorScan_spill iterative triv spill graph
                ksTriv ksSpill kksCoalesce

-- Select:
-- we couldn't find any triv nodes or things to freeze or coalesce,
--      and the graph isn't empty yet.. We'll have to choose a spill
--      candidate and leave it uncolored.
--
colorScan_spill
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Ord k,       Eq cls
           , Outputable k, Outputable cls)
        => Bool
        -> Triv k cls color
        -> (Graph k cls color -> k)
        -> Graph k cls color
        -> [k]
        -> [k]
        -> [(k, k)]
        -> ([k], [k], [(k, k)])

colorScan_spill iterative triv spill graph
        ksTriv ksSpill kksCoalesce

 = let  kSpill          = spill graph
        Just graph'     = delNode kSpill graph
   in   colorScan_spin iterative triv spill graph'
                ksTriv (kSpill : ksSpill) kksCoalesce

-- | Try to assign a color to all these nodes.

assignColors
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Outputable cls)
        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
        -> Graph k cls color            -- ^ the graph
        -> [k]                          -- ^ nodes to assign a color to.
        -> ( Graph k cls color          -- the colored graph
           , [k])                       -- the nodes that didn't color.

assignColors colors graph ks
        = assignColors' colors graph [] ks

 where  assignColors' _ graph prob []
                = (graph, prob)

        assignColors' colors graph prob (k:ks)
         = case assignColor colors k graph of

                -- couldn't color this node
                Nothing         -> assignColors' colors graph (k : prob) ks

                -- this node colored ok, so do the rest
                Just graph'     -> assignColors' colors graph' prob ks

        assignColor colors u graph
                | Just c        <- selectColor colors graph u
                = Just (setColor u c graph)

                | otherwise
                = Nothing

-- | Select a color for a certain node
--      taking into account preferences, neighbors and exclusions.
--      returns Nothing if no color can be assigned to this node.
--
selectColor
        :: ( Uniquable k, Uniquable cls, Uniquable color
           , Outputable cls)
        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
        -> Graph k cls color            -- ^ the graph
        -> k                            -- ^ key of the node to select a color for.
        -> Maybe color

selectColor colors graph u
 = let  -- lookup the node
        Just node       = lookupNode graph u

        -- lookup the available colors for the class of this node.
        colors_avail
         = case lookupUFM colors (nodeClass node) of
                Nothing -> pprPanic "selectColor: no colors available for class " (ppr (nodeClass node))
                Just cs -> cs

        -- find colors we can't use because they're already being used
        --      by a node that conflicts with this one.
        Just nsConflicts
                        = sequence
                        $ map (lookupNode graph)
                        $ nonDetEltsUniqSet
                        $ nodeConflicts node
                        -- See Note [Unique Determinism and code generation]

        colors_conflict = mkUniqSet
                        $ catMaybes
                        $ map nodeColor nsConflicts

        -- the prefs of our neighbors
        colors_neighbor_prefs
                        = mkUniqSet
                        $ concat $ map nodePreference nsConflicts

        -- colors that are still valid for us
        colors_ok_ex    = minusUniqSet colors_avail (nodeExclusions node)
        colors_ok       = minusUniqSet colors_ok_ex colors_conflict

        -- the colors that we prefer, and are still ok
        colors_ok_pref  = intersectUniqSets
                                (mkUniqSet $ nodePreference node) colors_ok

        -- the colors that we could choose while being nice to our neighbors
        colors_ok_nice  = minusUniqSet
                                colors_ok colors_neighbor_prefs

        -- the best of all possible worlds..
        colors_ok_pref_nice
                        = intersectUniqSets
                                colors_ok_nice colors_ok_pref

        -- make the decision
        chooseColor

                -- everyone is happy, yay!
                | not $ isEmptyUniqSet colors_ok_pref_nice
                , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref_nice)
                                        (nodePreference node)
                = Just c

                -- we've got one of our preferences
                | not $ isEmptyUniqSet colors_ok_pref
                , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref)
                                        (nodePreference node)
                = Just c

                -- it wasn't a preference, but it was still ok
                | not $ isEmptyUniqSet colors_ok
                , c : _         <- nonDetEltsUniqSet colors_ok
                -- See Note [Unique Determinism and code generation]
                = Just c

                -- no colors were available for us this time.
                --      looks like we're going around the loop again..
                | otherwise
                = Nothing

   in   chooseColor


-- | Basic operations on graphs.
--

module GraphOps (
        addNode,        delNode,        getNode,       lookupNode,     modNode,
        size,
        union,
        addConflict,    delConflict,    addConflicts,
        addCoalesce,    delCoalesce,
        addExclusion,   addExclusions,
        addPreference,
        coalesceNodes,  coalesceGraph,
        freezeNode,     freezeOneInGraph, freezeAllInGraph,
        scanGraph,
        setColor,
        validateGraph,
        slurpNodeConflictCount
)
where

import GhcPrelude

import GraphBase

import Outputable
import Unique
import UniqSet
import UniqFM

import Data.List        hiding (union)
import Data.Maybe

-- | Lookup a node from the graph.
lookupNode
        :: Uniquable k
        => Graph k cls color
        -> k -> Maybe (Node  k cls color)

lookupNode graph k
        = lookupUFM (graphMap graph) k

-- | Get a node from the graph, throwing an error if it's not there
getNode
        :: Uniquable k
        => Graph k cls color
        -> k -> Node k cls color

getNode graph k
 = case lookupUFM (graphMap graph) k of
        Just node       -> node
        Nothing         -> panic "ColorOps.getNode: not found"

-- | Add a node to the graph, linking up its edges
addNode :: Uniquable k
        => k -> Node k cls color
        -> Graph k cls color -> Graph k cls color

addNode k node graph
 = let
        -- add back conflict edges from other nodes to this one
        map_conflict =
          nonDetFoldUniqSet
            -- It's OK to use nonDetFoldUFM here because the
            -- operation is commutative
            (adjustUFM_C (\n -> n { nodeConflicts =
                                      addOneToUniqSet (nodeConflicts n) k}))
            (graphMap graph)
            (nodeConflicts node)

        -- add back coalesce edges from other nodes to this one
        map_coalesce =
          nonDetFoldUniqSet
            -- It's OK to use nonDetFoldUFM here because the
            -- operation is commutative
            (adjustUFM_C (\n -> n { nodeCoalesce =
                                      addOneToUniqSet (nodeCoalesce n) k}))
            map_conflict
            (nodeCoalesce node)

  in    graph
        { graphMap      = addToUFM map_coalesce k node}

-- | Delete a node and all its edges from the graph.
delNode :: (Uniquable k)
        => k -> Graph k cls color -> Maybe (Graph k cls color)

delNode k graph
        | Just node     <- lookupNode graph k
        = let   -- delete conflict edges from other nodes to this one.
                graph1  = foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
                        $ nonDetEltsUniqSet (nodeConflicts node)

                -- delete coalesce edge from other nodes to this one.
                graph2  = foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
                        $ nonDetEltsUniqSet (nodeCoalesce node)
                        -- See Note [Unique Determinism and code generation]

                -- delete the node
                graph3  = graphMapModify (\fm -> delFromUFM fm k) graph2

          in    Just graph3

        | otherwise
        = Nothing

-- | Modify a node in the graph.
--      returns Nothing if the node isn't present.
--
modNode :: Uniquable k
        => (Node k cls color -> Node k cls color)
        -> k -> Graph k cls color -> Maybe (Graph k cls color)

modNode f k graph
 = case lookupNode graph k of
        Just Node{}
         -> Just
         $  graphMapModify
                 (\fm   -> let  Just node       = lookupUFM fm k
                                node'           = f node
                           in   addToUFM fm k node')
                graph

        Nothing -> Nothing

-- | Get the size of the graph, O(n)
size    :: Graph k cls color -> Int

size graph
        = sizeUFM $ graphMap graph

-- | Union two graphs together.
union   :: Graph k cls color -> Graph k cls color -> Graph k cls color

union   graph1 graph2
        = Graph
        { graphMap              = plusUFM (graphMap graph1) (graphMap graph2) }

-- | Add a conflict between nodes to the graph, creating the nodes required.
--      Conflicts are virtual regs which need to be colored differently.
addConflict
        :: Uniquable k
        => (k, cls) -> (k, cls)
        -> Graph k cls color -> Graph k cls color

addConflict (u1, c1) (u2, c2)
 = let  addNeighbor u c u'
                = adjustWithDefaultUFM
                        (\node -> node { nodeConflicts = addOneToUniqSet (nodeConflicts node) u' })
                        (newNode u c)  { nodeConflicts = unitUniqSet u' }
                        u

   in   graphMapModify
        ( addNeighbor u1 c1 u2
        . addNeighbor u2 c2 u1)

-- | Delete a conflict edge. k1 -> k2
--      returns Nothing if the node isn't in the graph
delConflict
        :: Uniquable k
        => k -> k
        -> Graph k cls color -> Maybe (Graph k cls color)

delConflict k1 k2
        = modNode
                (\node -> node { nodeConflicts = delOneFromUniqSet (nodeConflicts node) k2 })
                k1

-- | Add some conflicts to the graph, creating nodes if required.
--      All the nodes in the set are taken to conflict with each other.
addConflicts
        :: Uniquable k
        => UniqSet k -> (k -> cls)
        -> Graph k cls color -> Graph k cls color

addConflicts conflicts getClass

        -- just a single node, but no conflicts, create the node anyway.
        | (u : [])      <- nonDetEltsUniqSet conflicts
        = graphMapModify
        $ adjustWithDefaultUFM
                id
                (newNode u (getClass u))
                u

        | otherwise
        = graphMapModify
        $ \fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
                $ nonDetEltsUniqSet conflicts
                -- See Note [Unique Determinism and code generation]

addConflictSet1 :: Uniquable k
                => k -> (k -> cls) -> UniqSet k
                -> UniqFM (Node k cls color)
                -> UniqFM (Node k cls color)
addConflictSet1 u getClass set
 = case delOneFromUniqSet set u of
    set' -> adjustWithDefaultUFM
                (\node -> node                  { nodeConflicts = unionUniqSets set' (nodeConflicts node) } )
                (newNode u (getClass u))        { nodeConflicts = set' }
                u

-- | Add an exclusion to the graph, creating nodes if required.
--      These are extra colors that the node cannot use.
addExclusion
        :: (Uniquable k, Uniquable color)
        => k -> (k -> cls) -> color
        -> Graph k cls color -> Graph k cls color

addExclusion u getClass color
        = graphMapModify
        $ adjustWithDefaultUFM
                (\node -> node                  { nodeExclusions = addOneToUniqSet (nodeExclusions node) color })
                (newNode u (getClass u))        { nodeExclusions = unitUniqSet color }
                u

addExclusions
        :: (Uniquable k, Uniquable color)
        => k -> (k -> cls) -> [color]
        -> Graph k cls color -> Graph k cls color

addExclusions u getClass colors graph
        = foldr (addExclusion u getClass) graph colors

-- | Add a coalescence edge to the graph, creating nodes if requried.
--      It is considered adventageous to assign the same color to nodes in a coalesence.
addCoalesce
        :: Uniquable k
        => (k, cls) -> (k, cls)
        -> Graph k cls color -> Graph k cls color

addCoalesce (u1, c1) (u2, c2)
 = let  addCoalesce u c u'
         =      adjustWithDefaultUFM
                        (\node -> node { nodeCoalesce = addOneToUniqSet (nodeCoalesce node) u' })
                        (newNode u c)  { nodeCoalesce = unitUniqSet u' }
                        u

   in   graphMapModify
        ( addCoalesce u1 c1 u2
        . addCoalesce u2 c2 u1)

-- | Delete a coalescence edge (k1 -> k2) from the graph.
delCoalesce
        :: Uniquable k
        => k -> k
        -> Graph k cls color    -> Maybe (Graph k cls color)

delCoalesce k1 k2
        = modNode (\node -> node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k2 })
                k1

-- | Add a color preference to the graph, creating nodes if required.
--      The most recently added preference is the most prefered.
--      The algorithm tries to assign a node it's prefered color if possible.
--
addPreference
        :: Uniquable k
        => (k, cls) -> color
        -> Graph k cls color -> Graph k cls color

addPreference (u, c) color
        = graphMapModify
        $ adjustWithDefaultUFM
                (\node -> node { nodePreference = color : (nodePreference node) })
                (newNode u c)  { nodePreference = [color] }
                u

-- | Do aggressive coalescing on this graph.
--      returns the new graph and the list of pairs of nodes that got coalesced together.
--      for each pair, the resulting node will have the least key and be second in the pair.
--
coalesceGraph
        :: (Uniquable k, Ord k, Eq cls, Outputable k)
        => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
                                --      less colorable (aggressive coalescing)
        -> Triv k cls color
        -> Graph k cls color
        -> ( Graph k cls color
           , [(k, k)])          -- pairs of nodes that were coalesced, in the order that the
                                --      coalescing was applied.

coalesceGraph aggressive triv graph
        = coalesceGraph' aggressive triv graph []

coalesceGraph'
        :: (Uniquable k, Ord k, Eq cls, Outputable k)
        => Bool
        -> Triv k cls color
        -> Graph k cls color
        -> [(k, k)]
        -> ( Graph k cls color
           , [(k, k)])
coalesceGraph' aggressive triv graph kkPairsAcc
 = let
        -- find all the nodes that have coalescence edges
        cNodes  = filter (\node -> not $ isEmptyUniqSet (nodeCoalesce node))
                $ nonDetEltsUFM $ graphMap graph
                -- See Note [Unique Determinism and code generation]

        -- build a list of pairs of keys for node's we'll try and coalesce
        --      every pair of nodes will appear twice in this list
        --      ie [(k1, k2), (k2, k1) ... ]
        --      This is ok, GrapOps.coalesceNodes handles this and it's convenient for
        --      build a list of what nodes get coalesced together for later on.
        --
        cList   = [ (nodeId node1, k2)
                        | node1 <- cNodes
                        , k2    <- nonDetEltsUniqSet $ nodeCoalesce node1 ]
                        -- See Note [Unique Determinism and code generation]

        -- do the coalescing, returning the new graph and a list of pairs of keys
        --      that got coalesced together.
        (graph', mPairs)
                = mapAccumL (coalesceNodes aggressive triv) graph cList

        -- keep running until there are no more coalesces can be found
   in   case catMaybes mPairs of
         []     -> (graph', reverse kkPairsAcc)
         pairs  -> coalesceGraph' aggressive triv graph' (reverse pairs ++ kkPairsAcc)

-- | Coalesce this pair of nodes unconditionally \/ aggressively.
--      The resulting node is the one with the least key.
--
--      returns: Just    the pair of keys if the nodes were coalesced
--                       the second element of the pair being the least one
--
--               Nothing if either of the nodes weren't in the graph

coalesceNodes
        :: (Uniquable k, Ord k, Eq cls)
        => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
                                --      less colorable (aggressive coalescing)
        -> Triv  k cls color
        -> Graph k cls color
        -> (k, k)               -- ^ keys of the nodes to be coalesced
        -> (Graph k cls color, Maybe (k, k))

coalesceNodes aggressive triv graph (k1, k2)
        | (kMin, kMax)  <- if k1 < k2
                                then (k1, k2)
                                else (k2, k1)

        -- the nodes being coalesced must be in the graph
        , Just nMin     <- lookupNode graph kMin
        , Just nMax     <- lookupNode graph kMax

        -- can't coalesce conflicting modes
        , not $ elementOfUniqSet kMin (nodeConflicts nMax)
        , not $ elementOfUniqSet kMax (nodeConflicts nMin)

        -- can't coalesce the same node
        , nodeId nMin /= nodeId nMax

        = coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax

        -- don't do the coalescing after all
        | otherwise
        = (graph, Nothing)

coalesceNodes_merge
        :: (Uniquable k, Eq cls)
        => Bool
        -> Triv  k cls color
        -> Graph k cls color
        -> k -> k
        -> Node k cls color
        -> Node k cls color
        -> (Graph k cls color, Maybe (k, k))

coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax

        -- sanity checks
        | nodeClass nMin /= nodeClass nMax
        = error "GraphOps.coalesceNodes: can't coalesce nodes of different classes."

        | not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
        = error "GraphOps.coalesceNodes: can't coalesce colored nodes."

        ---
        | otherwise
        = let
                -- the new node gets all the edges from its two components
                node    =
                 Node   { nodeId                = kMin
                        , nodeClass             = nodeClass nMin
                        , nodeColor             = Nothing

                        -- nodes don't conflict with themselves..
                        , nodeConflicts
                                = (unionUniqSets (nodeConflicts nMin) (nodeConflicts nMax))
                                        `delOneFromUniqSet` kMin
                                        `delOneFromUniqSet` kMax

                        , nodeExclusions        = unionUniqSets (nodeExclusions nMin) (nodeExclusions nMax)
                        , nodePreference        = nodePreference nMin ++ nodePreference nMax

                        -- nodes don't coalesce with themselves..
                        , nodeCoalesce
                                = (unionUniqSets (nodeCoalesce nMin) (nodeCoalesce nMax))
                                        `delOneFromUniqSet` kMin
                                        `delOneFromUniqSet` kMax
                        }

          in    coalesceNodes_check aggressive triv graph kMin kMax node

coalesceNodes_check
        :: Uniquable k
        => Bool
        -> Triv  k cls color
        -> Graph k cls color
        -> k -> k
        -> Node k cls color
        -> (Graph k cls color, Maybe (k, k))

coalesceNodes_check aggressive triv graph kMin kMax node

        -- Unless we're coalescing aggressively, if the result node is not trivially
        --      colorable then don't do the coalescing.
        | not aggressive
        , not $ triv (nodeClass node) (nodeConflicts node) (nodeExclusions node)
        = (graph, Nothing)

        | otherwise
        = let -- delete the old nodes from the graph and add the new one
                Just graph1     = delNode kMax graph
                Just graph2     = delNode kMin graph1
                graph3          = addNode kMin node graph2

          in    (graph3, Just (kMax, kMin))

-- | Freeze a node
--      This is for the iterative coalescer.
--      By freezing a node we give up on ever coalescing it.
--      Move all its coalesce edges into the frozen set - and update
--      back edges from other nodes.
--
freezeNode
        :: Uniquable k
        => k                    -- ^ key of the node to freeze
        -> Graph k cls color    -- ^ the graph
        -> Graph k cls color    -- ^ graph with that node frozen

freezeNode k
  = graphMapModify
  $ \fm ->
    let -- freeze all the edges in the node to be frozen
        Just node = lookupUFM fm k
        node'   = node
                { nodeCoalesce          = emptyUniqSet }

        fm1     = addToUFM fm k node'

        -- update back edges pointing to this node
        freezeEdge k node
         = if elementOfUniqSet k (nodeCoalesce node)
                then node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k }
                else node       -- panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
                                -- If the edge isn't actually in the coelesce set then just ignore it.

        fm2     = nonDetFoldUniqSet (adjustUFM_C (freezeEdge k)) fm1
                    -- It's OK to use nonDetFoldUFM here because the operation
                    -- is commutative
                        $ nodeCoalesce node

    in  fm2

-- | Freeze one node in the graph
--      This if for the iterative coalescer.
--      Look for a move related node of low degree and freeze it.
--
--      We probably don't need to scan the whole graph looking for the node of absolute
--      lowest degree. Just sample the first few and choose the one with the lowest
--      degree out of those. Also, we don't make any distinction between conflicts of different
--      classes.. this is just a heuristic, after all.
--
--      IDEA:   freezing a node might free it up for Simplify.. would be good to check for triv
--              right here, and add it to a worklist if known triv\/non-move nodes.
--
freezeOneInGraph
        :: (Uniquable k)
        => Graph k cls color
        -> ( Graph k cls color          -- the new graph
           , Bool )                     -- whether we found a node to freeze

freezeOneInGraph graph
 = let  compareNodeDegree n1 n2
                = compare (sizeUniqSet $ nodeConflicts n1) (sizeUniqSet $ nodeConflicts n2)

        candidates
                = sortBy compareNodeDegree
                $ take 5        -- 5 isn't special, it's just a small number.
                $ scanGraph (\node -> not $ isEmptyUniqSet (nodeCoalesce node)) graph

   in   case candidates of

         -- there wasn't anything available to freeze
         []     -> (graph, False)

         -- we found something to freeze
         (n : _)
          -> ( freezeNode (nodeId n) graph
             , True)

-- | Freeze all the nodes in the graph
--      for debugging the iterative allocator.
--
freezeAllInGraph
        :: (Uniquable k)
        => Graph k cls color
        -> Graph k cls color

freezeAllInGraph graph
        = foldr freezeNode graph
                $ map nodeId
                $ nonDetEltsUFM $ graphMap graph
                -- See Note [Unique Determinism and code generation]

-- | Find all the nodes in the graph that meet some criteria
--
scanGraph
        :: (Node k cls color -> Bool)
        -> Graph k cls color
        -> [Node k cls color]

scanGraph match graph
        = filter match $ nonDetEltsUFM $ graphMap graph
          -- See Note [Unique Determinism and code generation]

-- | validate the internal structure of a graph
--      all its edges should point to valid nodes
--      If they don't then throw an error
--
validateGraph
        :: (Uniquable k, Outputable k, Eq color)
        => SDoc                         -- ^ extra debugging info to display on error
        -> Bool                         -- ^ whether this graph is supposed to be colored.
        -> Graph k cls color            -- ^ graph to validate
        -> Graph k cls color            -- ^ validated graph

validateGraph doc isColored graph

        -- Check that all edges point to valid nodes.
        | edges         <- unionManyUniqSets
                                (  (map nodeConflicts       $ nonDetEltsUFM $ graphMap graph)
                                ++ (map nodeCoalesce        $ nonDetEltsUFM $ graphMap graph))

        , nodes         <- mkUniqSet $ map nodeId $ nonDetEltsUFM $ graphMap graph
        , badEdges      <- minusUniqSet edges nodes
        , not $ isEmptyUniqSet badEdges
        = pprPanic "GraphOps.validateGraph"
                (  text "Graph has edges that point to non-existent nodes"
                $$ text "  bad edges: " <> pprUFM (getUniqSet badEdges) (vcat . map ppr)
                $$ doc )

        -- Check that no conflicting nodes have the same color
        | badNodes      <- filter (not . (checkNode graph))
                        $ nonDetEltsUFM $ graphMap graph
                           -- See Note [Unique Determinism and code generation]
        , not $ null badNodes
        = pprPanic "GraphOps.validateGraph"
                (  text "Node has same color as one of it's conflicts"
                $$ text "  bad nodes: " <> hcat (map (ppr . nodeId) badNodes)
                $$ doc)

        -- If this is supposed to be a colored graph,
        --      check that all nodes have a color.
        | isColored
        , badNodes      <- filter (\n -> isNothing $ nodeColor n)
                        $  nonDetEltsUFM $ graphMap graph
        , not $ null badNodes
        = pprPanic "GraphOps.validateGraph"
                (  text "Supposably colored graph has uncolored nodes."
                $$ text "  uncolored nodes: " <> hcat (map (ppr . nodeId) badNodes)
                $$ doc )

        -- graph looks ok
        | otherwise
        = graph

-- | If this node is colored, check that all the nodes which
--      conflict with it have different colors.
checkNode
        :: (Uniquable k, Eq color)
        => Graph k cls color
        -> Node  k cls color
        -> Bool                 -- ^ True if this node is ok

checkNode graph node
        | Just color            <- nodeColor node
        , Just neighbors        <- sequence $ map (lookupNode graph)
                                $  nonDetEltsUniqSet $ nodeConflicts node
            -- See Note [Unique Determinism and code generation]

        , neighbourColors       <- catMaybes $ map nodeColor neighbors
        , elem color neighbourColors
        = False

        | otherwise
        = True

-- | Slurp out a map of how many nodes had a certain number of conflict neighbours

slurpNodeConflictCount
        :: Graph k cls color
        -> UniqFM (Int, Int)    -- ^ (conflict neighbours, num nodes with that many conflicts)

slurpNodeConflictCount graph
        = addListToUFM_C
                (\(c1, n1) (_, n2) -> (c1, n1 + n2))
                emptyUFM
        $ map   (\node
                  -> let count  = sizeUniqSet $ nodeConflicts node
                     in  (count, (count, 1)))
        $ nonDetEltsUFM
        -- See Note [Unique Determinism and code generation]
        $ graphMap graph

-- | Set the color of a certain node
setColor
        :: Uniquable k
        => k -> color
        -> Graph k cls color -> Graph k cls color

setColor u color
        = graphMapModify
        $ adjustUFM_C
                (\n -> n { nodeColor = Just color })
                u

{-# INLINE adjustWithDefaultUFM #-}
adjustWithDefaultUFM
        :: Uniquable k
        => (a -> a) -> a -> k
        -> UniqFM a -> UniqFM a

adjustWithDefaultUFM f def k map
        = addToUFM_C
                (\old _ -> f old)
                map
                k def

-- Argument order different from UniqFM's adjustUFM
{-# INLINE adjustUFM_C #-}
adjustUFM_C
        :: Uniquable k
        => (a -> a)
        -> k -> UniqFM a -> UniqFM a

adjustUFM_C f k map
 = case lookupUFM map k of
        Nothing -> map
        Just a  -> addToUFM map k (f a)