{-# LANGUAGE GADTs, TypeFamilies, FlexibleContexts, FlexibleInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

----------------------------------------------------------------------------
--
-- Pretty-printing of Cmm as (a superset of) C--
--
-- (c) The University of Glasgow 2004-2006
--
-----------------------------------------------------------------------------
--
-- This is where we walk over CmmNode emitting an external representation,
-- suitable for parsing, in a syntax strongly reminiscent of C--. This
-- is the "External Core" for the Cmm layer.
--
-- As such, this should be a well-defined syntax: we want it to look nice.
-- Thus, we try wherever possible to use syntax defined in [1],
-- "The C-- Reference Manual", http://www.cs.tufts.edu/~nr/c--/index.html. We
-- differ slightly, in some cases. For one, we use I8 .. I64 for types, rather
-- than C--'s bits8 .. bits64.
--
-- We try to ensure that all information available in the abstract
-- syntax is reproduced, or reproducible, in the concrete syntax.
-- Data that is not in printed out can be reconstructed according to
-- conventions used in the pretty printer. There are at least two such
-- cases:
--      1) if a value has wordRep type, the type is not appended in the
--      output.
--      2) MachOps that operate over wordRep type are printed in a
--      C-style, rather than as their internal MachRep name.
--
-- These conventions produce much more readable Cmm output.
--
-- A useful example pass over Cmm is in nativeGen/MachCodeGen.hs

module PprCmm
  ( module PprCmmDecl
  , module PprCmmExpr
  )
where

import GhcPrelude hiding (succ)

import BlockId ()
import CLabel
import Cmm
import CmmUtils
import CmmSwitch
import DynFlags
import FastString
import Outputable
import PprCmmDecl
import PprCmmExpr
import Util
import PprCore ()

import BasicTypes
import Hoopl.Block
import Hoopl.Graph

-------------------------------------------------
-- Outputable instances

instance Outputable CmmStackInfo where
    ppr = pprStackInfo

instance Outputable CmmTopInfo where
    ppr = pprTopInfo

instance Outputable (CmmNode e x) where
    ppr = pprNode

instance Outputable Convention where
    ppr = pprConvention

instance Outputable ForeignConvention where
    ppr = pprForeignConvention

instance Outputable ForeignTarget where
    ppr = pprForeignTarget

instance Outputable CmmReturnInfo where
    ppr = pprReturnInfo

instance Outputable (Block CmmNode C C) where
    ppr = pprBlock
instance Outputable (Block CmmNode C O) where
    ppr = pprBlock
instance Outputable (Block CmmNode O C) where
    ppr = pprBlock
instance Outputable (Block CmmNode O O) where
    ppr = pprBlock

instance Outputable (Graph CmmNode e x) where
    ppr = pprGraph

instance Outputable CmmGraph where
    ppr = pprCmmGraph

----------------------------------------------------------
-- Outputting types Cmm contains

pprStackInfo :: CmmStackInfo -> SDoc
pprStackInfo (StackInfo {arg_space=arg_space, updfr_space=updfr_space}) =
  text "arg_space: " <> ppr arg_space <+>
  text "updfr_space: " <> ppr updfr_space

pprTopInfo :: CmmTopInfo -> SDoc
pprTopInfo (TopInfo {info_tbls=info_tbl, stack_info=stack_info}) =
  vcat [text "info_tbls: " <> ppr info_tbl,
        text "stack_info: " <> ppr stack_info]

----------------------------------------------------------
-- Outputting blocks and graphs

pprBlock :: IndexedCO x SDoc SDoc ~ SDoc
         => Block CmmNode e x -> IndexedCO e SDoc SDoc
pprBlock block
    = foldBlockNodesB3 ( ($$) . ppr
                       , ($$) . (nest 4) . ppr
                       , ($$) . (nest 4) . ppr
                       )
                       block
                       empty

pprGraph :: Graph CmmNode e x -> SDoc
pprGraph GNil = empty
pprGraph (GUnit block) = ppr block
pprGraph (GMany entry body exit)
   = text "{"
  $$ nest 2 (pprMaybeO entry $$ (vcat $ map ppr $ bodyToBlockList body) $$ pprMaybeO exit)
  $$ text "}"
  where pprMaybeO :: Outputable (Block CmmNode e x)
                  => MaybeO ex (Block CmmNode e x) -> SDoc
        pprMaybeO NothingO = empty
        pprMaybeO (JustO block) = ppr block

pprCmmGraph :: CmmGraph -> SDoc
pprCmmGraph g
   = text "{" <> text "offset"
  $$ nest 2 (vcat $ map ppr blocks)
  $$ text "}"
  where blocks = revPostorder g
    -- revPostorder has the side-effect of discarding unreachable code,
    -- so pretty-printed Cmm will omit any unreachable blocks.  This can
    -- sometimes be confusing.

---------------------------------------------
-- Outputting CmmNode and types which it contains

pprConvention :: Convention -> SDoc
pprConvention (NativeNodeCall   {}) = text "<native-node-call-convention>"
pprConvention (NativeDirectCall {}) = text "<native-direct-call-convention>"
pprConvention (NativeReturn {})     = text "<native-ret-convention>"
pprConvention  Slow                 = text "<slow-convention>"
pprConvention  GC                   = text "<gc-convention>"

pprForeignConvention :: ForeignConvention -> SDoc
pprForeignConvention (ForeignConvention c args res ret) =
          doubleQuotes (ppr c) <+> text "arg hints: " <+> ppr args <+> text " result hints: " <+> ppr res <+> ppr ret

pprReturnInfo :: CmmReturnInfo -> SDoc
pprReturnInfo CmmMayReturn = empty
pprReturnInfo CmmNeverReturns = text "never returns"

pprForeignTarget :: ForeignTarget -> SDoc
pprForeignTarget (ForeignTarget fn c) = ppr c <+> ppr_target fn
  where
        ppr_target :: CmmExpr -> SDoc
        ppr_target t@(CmmLit _) = ppr t
        ppr_target fn'          = parens (ppr fn')

pprForeignTarget (PrimTarget op)
 -- HACK: We're just using a ForeignLabel to get this printed, the label
 --       might not really be foreign.
 = ppr
               (CmmLabel (mkForeignLabel
                         (mkFastString (show op))
                         Nothing ForeignLabelInThisPackage IsFunction))

pprNode :: CmmNode e x -> SDoc
pprNode node = pp_node <+> pp_debug
  where
    pp_node :: SDoc
    pp_node = sdocWithDynFlags $ \dflags -> case node of
      -- label:
      CmmEntry id tscope -> lbl <> colon <+>
         (sdocWithDynFlags $ \dflags ->
           ppUnless (gopt Opt_SuppressTicks dflags) (text "//" <+> ppr tscope))
          where
            lbl = if gopt Opt_SuppressUniques dflags
                then text "_lbl_"
                else ppr id

      -- // text
      CmmComment s -> text "//" <+> ftext s

      -- //tick bla<...>
      CmmTick t -> ppUnless (gopt Opt_SuppressTicks dflags) $
                   text "//tick" <+> ppr t

      -- unwind reg = expr;
      CmmUnwind regs ->
          text "unwind "
          <> commafy (map (\(r,e) -> ppr r <+> char '=' <+> ppr e) regs) <> semi

      -- reg = expr;
      CmmAssign reg expr -> ppr reg <+> equals <+> ppr expr <> semi

      -- rep[lv] = expr;
      CmmStore lv expr -> rep <> brackets(ppr lv) <+> equals <+> ppr expr <> semi
          where
            rep = sdocWithDynFlags $ \dflags ->
                  ppr ( cmmExprType dflags expr )

      -- call "ccall" foo(x, y)[r1, r2];
      -- ToDo ppr volatile
      CmmUnsafeForeignCall target results args ->
          hsep [ ppUnless (null results) $
                    parens (commafy $ map ppr results) <+> equals,
                 text "call",
                 ppr target <> parens (commafy $ map ppr args) <> semi]

      -- goto label;
      CmmBranch ident -> text "goto" <+> ppr ident <> semi

      -- if (expr) goto t; else goto f;
      CmmCondBranch expr t f l ->
          hsep [ text "if"
               , parens(ppr expr)
               , case l of
                   Nothing -> empty
                   Just b -> parens (text "likely:" <+> ppr b)
               , text "goto"
               , ppr t <> semi
               , text "else goto"
               , ppr f <> semi
               ]

      CmmSwitch expr ids ->
          hang (hsep [ text "switch"
                     , range
                     , if isTrivialCmmExpr expr
                       then ppr expr
                       else parens (ppr expr)
                     , text "{"
                     ])
             4 (vcat (map ppCase cases) $$ def) $$ rbrace
          where
            (cases, mbdef) = switchTargetsFallThrough ids
            ppCase (is,l) = hsep
                            [ text "case"
                            , commafy $ map integer is
                            , text ": goto"
                            , ppr l <> semi
                            ]
            def | Just l <- mbdef = hsep
                            [ text "default:"
                            , braces (text "goto" <+> ppr l <> semi)
                            ]
                | otherwise = empty

            range = brackets $ hsep [integer lo, text "..", integer hi]
              where (lo,hi) = switchTargetsRange ids

      CmmCall tgt k regs out res updfr_off ->
          hcat [ text "call", space
               , pprFun tgt, parens (interpp'SP regs), space
               , returns <+>
                 text "args: " <> ppr out <> comma <+>
                 text "res: " <> ppr res <> comma <+>
                 text "upd: " <> ppr updfr_off
               , semi ]
          where pprFun f@(CmmLit _) = ppr f
                pprFun f = parens (ppr f)

                returns
                  | Just r <- k = text "returns to" <+> ppr r <> comma
                  | otherwise   = empty

      CmmForeignCall {tgt=t, res=rs, args=as, succ=s, ret_args=a, ret_off=u, intrbl=i} ->
          hcat $ if i then [text "interruptible", space] else [] ++
               [ text "foreign call", space
               , ppr t, text "(...)", space
               , text "returns to" <+> ppr s
                    <+> text "args:" <+> parens (ppr as)
                    <+> text "ress:" <+> parens (ppr rs)
               , text "ret_args:" <+> ppr a
               , text "ret_off:" <+> ppr u
               , semi ]

    pp_debug :: SDoc
    pp_debug =
      if not debugIsOn then empty
      else case node of
             CmmEntry {}             -> empty -- Looks terrible with text "  // CmmEntry"
             CmmComment {}           -> empty -- Looks also terrible with text "  // CmmComment"
             CmmTick {}              -> empty
             CmmUnwind {}            -> text "  // CmmUnwind"
             CmmAssign {}            -> text "  // CmmAssign"
             CmmStore {}             -> text "  // CmmStore"
             CmmUnsafeForeignCall {} -> text "  // CmmUnsafeForeignCall"
             CmmBranch {}            -> text "  // CmmBranch"
             CmmCondBranch {}        -> text "  // CmmCondBranch"
             CmmSwitch {}            -> text "  // CmmSwitch"
             CmmCall {}              -> text "  // CmmCall"
             CmmForeignCall {}       -> text "  // CmmForeignCall"

    commafy :: [SDoc] -> SDoc
    commafy xs = hsep $ punctuate comma xs


----------------------------------------------------------------------------
--
-- Pretty-printing of common Cmm types
--
-- (c) The University of Glasgow 2004-2006
--
-----------------------------------------------------------------------------

--
-- This is where we walk over Cmm emitting an external representation,
-- suitable for parsing, in a syntax strongly reminiscent of C--. This
-- is the "External Core" for the Cmm layer.
--
-- As such, this should be a well-defined syntax: we want it to look nice.
-- Thus, we try wherever possible to use syntax defined in [1],
-- "The C-- Reference Manual", http://www.cs.tufts.edu/~nr/c--/index.html. We
-- differ slightly, in some cases. For one, we use I8 .. I64 for types, rather
-- than C--'s bits8 .. bits64.
--
-- We try to ensure that all information available in the abstract
-- syntax is reproduced, or reproducible, in the concrete syntax.
-- Data that is not in printed out can be reconstructed according to
-- conventions used in the pretty printer. There are at least two such
-- cases:
--      1) if a value has wordRep type, the type is not appended in the
--      output.
--      2) MachOps that operate over wordRep type are printed in a
--      C-style, rather than as their internal MachRep name.
--
-- These conventions produce much more readable Cmm output.
--
-- A useful example pass over Cmm is in nativeGen/MachCodeGen.hs
--

{-# OPTIONS_GHC -fno-warn-orphans #-}
module PprCmmDecl
    ( writeCmms, pprCmms, pprCmmGroup, pprSection, pprStatic
    )
where

import GhcPrelude

import PprCmmExpr
import Cmm

import DynFlags
import Outputable
import FastString

import Data.List
import System.IO

import qualified Data.ByteString as BS

pprCmms :: (Outputable info, Outputable g)
        => [GenCmmGroup CmmStatics info g] -> SDoc
pprCmms cmms = pprCode CStyle (vcat (intersperse separator $ map ppr cmms))
        where
          separator = space $$ text "-------------------" $$ space

writeCmms :: (Outputable info, Outputable g)
          => DynFlags -> Handle -> [GenCmmGroup CmmStatics info g] -> IO ()
writeCmms dflags handle cmms = printForC dflags handle (pprCmms cmms)

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

instance (Outputable d, Outputable info, Outputable i)
      => Outputable (GenCmmDecl d info i) where
    ppr t = pprTop t

instance Outputable CmmStatics where
    ppr = pprStatics

instance Outputable CmmStatic where
    ppr = pprStatic

instance Outputable CmmInfoTable where
    ppr = pprInfoTable

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

pprCmmGroup :: (Outputable d, Outputable info, Outputable g)
            => GenCmmGroup d info g -> SDoc
pprCmmGroup tops
    = vcat $ intersperse blankLine $ map pprTop tops

-- --------------------------------------------------------------------------
-- Top level `procedure' blocks.
--
pprTop :: (Outputable d, Outputable info, Outputable i)
       => GenCmmDecl d info i -> SDoc

pprTop (CmmProc info lbl live graph)

  = vcat [ ppr lbl <> lparen <> rparen <+> text "// " <+> ppr live
         , nest 8 $ lbrace <+> ppr info $$ rbrace
         , nest 4 $ ppr graph
         , rbrace ]

-- --------------------------------------------------------------------------
-- We follow [1], 4.5
--
--      section "data" { ... }
--
pprTop (CmmData section ds) =
    (hang (pprSection section <+> lbrace) 4 (ppr ds))
    $$ rbrace

-- --------------------------------------------------------------------------
-- Info tables.

pprInfoTable :: CmmInfoTable -> SDoc
pprInfoTable (CmmInfoTable { cit_lbl = lbl, cit_rep = rep
                           , cit_prof = prof_info
                           , cit_srt = srt })
  = vcat [ text "label: " <> ppr lbl
         , text "rep: " <> ppr rep
         , case prof_info of
             NoProfilingInfo -> empty
             ProfilingInfo ct cd ->
               vcat [ text "type: " <> text (show (BS.unpack ct))
                    , text "desc: " <> text (show (BS.unpack cd)) ]
         , text "srt: " <> ppr srt ]

instance Outputable ForeignHint where
  ppr NoHint     = empty
  ppr SignedHint = quotes(text "signed")
--  ppr AddrHint   = quotes(text "address")
-- Temp Jan08
  ppr AddrHint   = (text "PtrHint")

-- --------------------------------------------------------------------------
-- Static data.
--      Strings are printed as C strings, and we print them as I8[],
--      following C--
--
pprStatics :: CmmStatics -> SDoc
pprStatics (Statics lbl ds) = vcat ((ppr lbl <> colon) : map ppr ds)

pprStatic :: CmmStatic -> SDoc
pprStatic s = case s of
    CmmStaticLit lit   -> nest 4 $ text "const" <+> pprLit lit <> semi
    CmmUninitialised i -> nest 4 $ text "I8" <> brackets (int i)
    CmmString s'       -> nest 4 $ text "I8[]" <+> text (show s')

-- --------------------------------------------------------------------------
-- data sections
--
pprSection :: Section -> SDoc
pprSection (Section t suffix) =
  section <+> doubleQuotes (pprSectionType t <+> char '.' <+> ppr suffix)
  where
    section = text "section"

pprSectionType :: SectionType -> SDoc
pprSectionType s = doubleQuotes (ptext t)
 where
  t = case s of
    Text              -> sLit "text"
    Data              -> sLit "data"
    ReadOnlyData      -> sLit "readonly"
    ReadOnlyData16    -> sLit "readonly16"
    RelocatableReadOnlyData
                      -> sLit "relreadonly"
    UninitialisedData -> sLit "uninitialised"
    CString           -> sLit "cstring"
    OtherSection s'   -> sLit s' -- Not actually a literal though.


----------------------------------------------------------------------------
--
-- Pretty-printing of common Cmm types
--
-- (c) The University of Glasgow 2004-2006
--
-----------------------------------------------------------------------------

--
-- This is where we walk over Cmm emitting an external representation,
-- suitable for parsing, in a syntax strongly reminiscent of C--. This
-- is the "External Core" for the Cmm layer.
--
-- As such, this should be a well-defined syntax: we want it to look nice.
-- Thus, we try wherever possible to use syntax defined in [1],
-- "The C-- Reference Manual", http://www.cs.tufts.edu/~nr/c--/index.html. We
-- differ slightly, in some cases. For one, we use I8 .. I64 for types, rather
-- than C--'s bits8 .. bits64.
--
-- We try to ensure that all information available in the abstract
-- syntax is reproduced, or reproducible, in the concrete syntax.
-- Data that is not in printed out can be reconstructed according to
-- conventions used in the pretty printer. There are at least two such
-- cases:
--      1) if a value has wordRep type, the type is not appended in the
--      output.
--      2) MachOps that operate over wordRep type are printed in a
--      C-style, rather than as their internal MachRep name.
--
-- These conventions produce much more readable Cmm output.
--
-- A useful example pass over Cmm is in nativeGen/MachCodeGen.hs
--

{-# OPTIONS_GHC -fno-warn-orphans #-}
module PprCmmExpr
    ( pprExpr, pprLit
    )
where

import GhcPrelude

import CmmExpr

import Outputable
import DynFlags

import Data.Maybe
import Numeric ( fromRat )

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

instance Outputable CmmExpr where
    ppr e = pprExpr e

instance Outputable CmmReg where
    ppr e = pprReg e

instance Outputable CmmLit where
    ppr l = pprLit l

instance Outputable LocalReg where
    ppr e = pprLocalReg e

instance Outputable Area where
    ppr e = pprArea e

instance Outputable GlobalReg where
    ppr e = pprGlobalReg e

-- --------------------------------------------------------------------------
-- Expressions
--

pprExpr :: CmmExpr -> SDoc
pprExpr e
    = sdocWithDynFlags $ \dflags ->
      case e of
        CmmRegOff reg i ->
                pprExpr (CmmMachOp (MO_Add rep)
                           [CmmReg reg, CmmLit (CmmInt (fromIntegral i) rep)])
                where rep = typeWidth (cmmRegType dflags reg)
        CmmLit lit -> pprLit lit
        _other     -> pprExpr1 e

-- Here's the precedence table from CmmParse.y:
-- %nonassoc '>=' '>' '<=' '<' '!=' '=='
-- %left '|'
-- %left '^'
-- %left '&'
-- %left '>>' '<<'
-- %left '-' '+'
-- %left '/' '*' '%'
-- %right '~'

-- We just cope with the common operators for now, the rest will get
-- a default conservative behaviour.

-- %nonassoc '>=' '>' '<=' '<' '!=' '=='
pprExpr1, pprExpr7, pprExpr8 :: CmmExpr -> SDoc
pprExpr1 (CmmMachOp op [x,y]) | Just doc <- infixMachOp1 op
   = pprExpr7 x <+> doc <+> pprExpr7 y
pprExpr1 e = pprExpr7 e

infixMachOp1, infixMachOp7, infixMachOp8 :: MachOp -> Maybe SDoc

infixMachOp1 (MO_Eq     _) = Just (text "==")
infixMachOp1 (MO_Ne     _) = Just (text "!=")
infixMachOp1 (MO_Shl    _) = Just (text "<<")
infixMachOp1 (MO_U_Shr  _) = Just (text ">>")
infixMachOp1 (MO_U_Ge   _) = Just (text ">=")
infixMachOp1 (MO_U_Le   _) = Just (text "<=")
infixMachOp1 (MO_U_Gt   _) = Just (char '>')
infixMachOp1 (MO_U_Lt   _) = Just (char '<')
infixMachOp1 _             = Nothing

-- %left '-' '+'
pprExpr7 (CmmMachOp (MO_Add rep1) [x, CmmLit (CmmInt i rep2)]) | i < 0
   = pprExpr7 (CmmMachOp (MO_Sub rep1) [x, CmmLit (CmmInt (negate i) rep2)])
pprExpr7 (CmmMachOp op [x,y]) | Just doc <- infixMachOp7 op
   = pprExpr7 x <+> doc <+> pprExpr8 y
pprExpr7 e = pprExpr8 e

infixMachOp7 (MO_Add _)  = Just (char '+')
infixMachOp7 (MO_Sub _)  = Just (char '-')
infixMachOp7 _           = Nothing

-- %left '/' '*' '%'
pprExpr8 (CmmMachOp op [x,y]) | Just doc <- infixMachOp8 op
   = pprExpr8 x <+> doc <+> pprExpr9 y
pprExpr8 e = pprExpr9 e

infixMachOp8 (MO_U_Quot _) = Just (char '/')
infixMachOp8 (MO_Mul _)    = Just (char '*')
infixMachOp8 (MO_U_Rem _)  = Just (char '%')
infixMachOp8 _             = Nothing

pprExpr9 :: CmmExpr -> SDoc
pprExpr9 e =
   case e of
        CmmLit    lit       -> pprLit1 lit
        CmmLoad   expr rep  -> ppr rep <> brackets (ppr expr)
        CmmReg    reg       -> ppr reg
        CmmRegOff  reg off  -> parens (ppr reg <+> char '+' <+> int off)
        CmmStackSlot a off  -> parens (ppr a   <+> char '+' <+> int off)
        CmmMachOp mop args  -> genMachOp mop args

genMachOp :: MachOp -> [CmmExpr] -> SDoc
genMachOp mop args
   | Just doc <- infixMachOp mop = case args of
        -- dyadic
        [x,y] -> pprExpr9 x <+> doc <+> pprExpr9 y

        -- unary
        [x]   -> doc <> pprExpr9 x

        _     -> pprTrace "PprCmm.genMachOp: machop with strange number of args"
                          (pprMachOp mop <+>
                            parens (hcat $ punctuate comma (map pprExpr args)))
                          empty

   | isJust (infixMachOp1 mop)
   || isJust (infixMachOp7 mop)
   || isJust (infixMachOp8 mop)  = parens (pprExpr (CmmMachOp mop args))

   | otherwise = char '%' <> ppr_op <> parens (commafy (map pprExpr args))
        where ppr_op = text (map (\c -> if c == ' ' then '_' else c)
                                 (show mop))
                -- replace spaces in (show mop) with underscores,

--
-- Unsigned ops on the word size of the machine get nice symbols.
-- All else get dumped in their ugly format.
--
infixMachOp :: MachOp -> Maybe SDoc
infixMachOp mop
        = case mop of
            MO_And    _ -> Just $ char '&'
            MO_Or     _ -> Just $ char '|'
            MO_Xor    _ -> Just $ char '^'
            MO_Not    _ -> Just $ char '~'
            MO_S_Neg  _ -> Just $ char '-' -- there is no unsigned neg :)
            _ -> Nothing

-- --------------------------------------------------------------------------
-- Literals.
--  To minimise line noise we adopt the convention that if the literal
--  has the natural machine word size, we do not append the type
--
pprLit :: CmmLit -> SDoc
pprLit lit = sdocWithDynFlags $ \dflags ->
             case lit of
    CmmInt i rep ->
        hcat [ (if i < 0 then parens else id)(integer i)
             , ppUnless (rep == wordWidth dflags) $
               space <> dcolon <+> ppr rep ]

    CmmFloat f rep     -> hsep [ double (fromRat f), dcolon, ppr rep ]
    CmmVec lits        -> char '<' <> commafy (map pprLit lits) <> char '>'
    CmmLabel clbl      -> ppr clbl
    CmmLabelOff clbl i -> ppr clbl <> ppr_offset i
    CmmLabelDiffOff clbl1 clbl2 i _ -> ppr clbl1 <> char '-'
                                  <> ppr clbl2 <> ppr_offset i
    CmmBlock id        -> ppr id
    CmmHighStackMark -> text "<highSp>"

pprLit1 :: CmmLit -> SDoc
pprLit1 lit@(CmmLabelOff {}) = parens (pprLit lit)
pprLit1 lit                  = pprLit lit

ppr_offset :: Int -> SDoc
ppr_offset i
    | i==0      = empty
    | i>=0      = char '+' <> int i
    | otherwise = char '-' <> int (-i)

-- --------------------------------------------------------------------------
-- Registers, whether local (temps) or global
--
pprReg :: CmmReg -> SDoc
pprReg r
    = case r of
        CmmLocal  local  -> pprLocalReg  local
        CmmGlobal global -> pprGlobalReg global

--
-- We only print the type of the local reg if it isn't wordRep
--
pprLocalReg :: LocalReg -> SDoc
pprLocalReg (LocalReg uniq rep) = sdocWithDynFlags $ \dflags ->
--   = ppr rep <> char '_' <> ppr uniq
-- Temp Jan08
    char '_' <> pprUnique dflags uniq <>
       (if isWord32 rep -- && not (isGcPtrType rep) -- Temp Jan08               -- sigh
                    then dcolon <> ptr <> ppr rep
                    else dcolon <> ptr <> ppr rep)
   where
     pprUnique dflags unique =
        if gopt Opt_SuppressUniques dflags
            then text "_locVar_"
            else ppr unique
     ptr = empty
         --if isGcPtrType rep
         --      then doubleQuotes (text "ptr")
         --      else empty

-- Stack areas
pprArea :: Area -> SDoc
pprArea Old        = text "old"
pprArea (Young id) = hcat [ text "young<", ppr id, text ">" ]

-- needs to be kept in syn with CmmExpr.hs.GlobalReg
--
pprGlobalReg :: GlobalReg -> SDoc
pprGlobalReg gr
    = case gr of
        VanillaReg n _ -> char 'R' <> int n
-- Temp Jan08
--        VanillaReg n VNonGcPtr -> char 'R' <> int n
--        VanillaReg n VGcPtr    -> char 'P' <> int n
        FloatReg   n   -> char 'F' <> int n
        DoubleReg  n   -> char 'D' <> int n
        LongReg    n   -> char 'L' <> int n
        XmmReg     n   -> text "XMM" <> int n
        YmmReg     n   -> text "YMM" <> int n
        ZmmReg     n   -> text "ZMM" <> int n
        Sp             -> text "Sp"
        SpLim          -> text "SpLim"
        Hp             -> text "Hp"
        HpLim          -> text "HpLim"
        MachSp         -> text "MachSp"
        UnwindReturnReg-> text "UnwindReturnReg"
        CCCS           -> text "CCCS"
        CurrentTSO     -> text "CurrentTSO"
        CurrentNursery -> text "CurrentNursery"
        HpAlloc        -> text "HpAlloc"
        EagerBlackholeInfo -> text "stg_EAGER_BLACKHOLE_info"
        GCEnter1       -> text "stg_gc_enter_1"
        GCFun          -> text "stg_gc_fun"
        BaseReg        -> text "BaseReg"
        PicBaseReg     -> text "PicBaseReg"

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

commafy :: [SDoc] -> SDoc
commafy xs = fsep $ punctuate comma xs


-- (c) The University of Glasgow 2006
-- (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
--
-- Storage manager representation of closures

{-# LANGUAGE CPP,GeneralizedNewtypeDeriving #-}

module SMRep (
        -- * Words and bytes
        WordOff, ByteOff,
        wordsToBytes, bytesToWordsRoundUp,
        roundUpToWords, roundUpTo,

        StgWord, fromStgWord, toStgWord,
        StgHalfWord, fromStgHalfWord, toStgHalfWord,
        hALF_WORD_SIZE, hALF_WORD_SIZE_IN_BITS,

        -- * Closure repesentation
        SMRep(..), -- CmmInfo sees the rep; no one else does
        IsStatic,
        ClosureTypeInfo(..), ArgDescr(..), Liveness,
        ConstrDescription,

        -- ** Construction
        mkHeapRep, blackHoleRep, indStaticRep, mkStackRep, mkRTSRep, arrPtrsRep,
        smallArrPtrsRep, arrWordsRep,

        -- ** Predicates
        isStaticRep, isConRep, isThunkRep, isFunRep, isStaticNoCafCon,
        isStackRep,

        -- ** Size-related things
        heapClosureSizeW,
        fixedHdrSizeW, arrWordsHdrSize, arrWordsHdrSizeW, arrPtrsHdrSize,
        arrPtrsHdrSizeW, profHdrSize, thunkHdrSize, nonHdrSize, nonHdrSizeW,
        smallArrPtrsHdrSize, smallArrPtrsHdrSizeW, hdrSize, hdrSizeW,
        fixedHdrSize,

        -- ** RTS closure types
        rtsClosureType, rET_SMALL, rET_BIG,
        aRG_GEN, aRG_GEN_BIG,

        -- ** Arrays
        card, cardRoundUp, cardTableSizeB, cardTableSizeW
    ) where

import GhcPrelude

import BasicTypes( ConTagZ )
import DynFlags
import Outputable
import Platform
import FastString

import Data.Word
import Data.Bits
import Data.ByteString (ByteString)

{-
************************************************************************
*                                                                      *
                Words and bytes
*                                                                      *
************************************************************************
-}

-- | Word offset, or word count
type WordOff = Int

-- | Byte offset, or byte count
type ByteOff = Int

-- | Round up the given byte count to the next byte count that's a
-- multiple of the machine's word size.
roundUpToWords :: DynFlags -> ByteOff -> ByteOff
roundUpToWords dflags n = roundUpTo n (wORD_SIZE dflags)

-- | Round up @base@ to a multiple of @size@.
roundUpTo :: ByteOff -> ByteOff -> ByteOff
roundUpTo base size = (base + (size - 1)) .&. (complement (size - 1))

-- | Convert the given number of words to a number of bytes.
--
-- This function morally has type @WordOff -> ByteOff@, but uses @Num
-- a@ to allow for overloading.
wordsToBytes :: Num a => DynFlags -> a -> a
wordsToBytes dflags n = fromIntegral (wORD_SIZE dflags) * n
{-# SPECIALIZE wordsToBytes :: DynFlags -> Int -> Int #-}
{-# SPECIALIZE wordsToBytes :: DynFlags -> Word -> Word #-}
{-# SPECIALIZE wordsToBytes :: DynFlags -> Integer -> Integer #-}

-- | First round the given byte count up to a multiple of the
-- machine's word size and then convert the result to words.
bytesToWordsRoundUp :: DynFlags -> ByteOff -> WordOff
bytesToWordsRoundUp dflags n = (n + word_size - 1) `quot` word_size
 where word_size = wORD_SIZE dflags
-- StgWord is a type representing an StgWord on the target platform.
-- A Word64 is large enough to hold a Word for either a 32bit or 64bit platform
newtype StgWord = StgWord Word64
    deriving (Eq, Bits)

fromStgWord :: StgWord -> Integer
fromStgWord (StgWord i) = toInteger i

toStgWord :: DynFlags -> Integer -> StgWord
toStgWord dflags i
    = case platformWordSize (targetPlatform dflags) of
      -- These conversions mean that things like toStgWord (-1)
      -- do the right thing
      4 -> StgWord (fromIntegral (fromInteger i :: Word32))
      8 -> StgWord (fromInteger i :: Word64)
      w -> panic ("toStgWord: Unknown platformWordSize: " ++ show w)

instance Outputable StgWord where
    ppr (StgWord i) = integer (toInteger i)

--

-- A Word32 is large enough to hold half a Word for either a 32bit or
-- 64bit platform
newtype StgHalfWord = StgHalfWord Word32
    deriving Eq

fromStgHalfWord :: StgHalfWord -> Integer
fromStgHalfWord (StgHalfWord w) = toInteger w

toStgHalfWord :: DynFlags -> Integer -> StgHalfWord
toStgHalfWord dflags i
    = case platformWordSize (targetPlatform dflags) of
      -- These conversions mean that things like toStgHalfWord (-1)
      -- do the right thing
      4 -> StgHalfWord (fromIntegral (fromInteger i :: Word16))
      8 -> StgHalfWord (fromInteger i :: Word32)
      w -> panic ("toStgHalfWord: Unknown platformWordSize: " ++ show w)

instance Outputable StgHalfWord where
    ppr (StgHalfWord w) = integer (toInteger w)

hALF_WORD_SIZE :: DynFlags -> ByteOff
hALF_WORD_SIZE dflags = platformWordSize (targetPlatform dflags) `shiftR` 1
hALF_WORD_SIZE_IN_BITS :: DynFlags -> Int
hALF_WORD_SIZE_IN_BITS dflags = platformWordSize (targetPlatform dflags) `shiftL` 2

{-
************************************************************************
*                                                                      *
\subsubsection[SMRep-datatype]{@SMRep@---storage manager representation}
*                                                                      *
************************************************************************
-}

-- | A description of the layout of a closure.  Corresponds directly
-- to the closure types in includes/rts/storage/ClosureTypes.h.
data SMRep
  = HeapRep              -- GC routines consult sizes in info tbl
        IsStatic
        !WordOff         --  # ptr words
        !WordOff         --  # non-ptr words INCLUDING SLOP (see mkHeapRep below)
        ClosureTypeInfo  -- type-specific info

  | ArrayPtrsRep
        !WordOff        -- # ptr words
        !WordOff        -- # card table words

  | SmallArrayPtrsRep
        !WordOff        -- # ptr words

  | ArrayWordsRep
        !WordOff        -- # bytes expressed in words, rounded up

  | StackRep            -- Stack frame (RET_SMALL or RET_BIG)
        Liveness

  | RTSRep              -- The RTS needs to declare info tables with specific
        Int             -- type tags, so this form lets us override the default
        SMRep           -- tag for an SMRep.

-- | True <=> This is a static closure.  Affects how we garbage-collect it.
-- Static closure have an extra static link field at the end.
-- Constructors do not have a static variant; see Note [static constructors]
type IsStatic = Bool

-- From an SMRep you can get to the closure type defined in
-- includes/rts/storage/ClosureTypes.h. Described by the function
-- rtsClosureType below.

data ClosureTypeInfo
  = Constr        ConTagZ ConstrDescription
  | Fun           FunArity ArgDescr
  | Thunk
  | ThunkSelector SelectorOffset
  | BlackHole
  | IndStatic

type ConstrDescription = ByteString -- result of dataConIdentity
type FunArity          = Int
type SelectorOffset    = Int

-------------------------
-- We represent liveness bitmaps as a Bitmap (whose internal
-- representation really is a bitmap).  These are pinned onto case return
-- vectors to indicate the state of the stack for the garbage collector.
--
-- In the compiled program, liveness bitmaps that fit inside a single
-- word (StgWord) are stored as a single word, while larger bitmaps are
-- stored as a pointer to an array of words.

type Liveness = [Bool]   -- One Bool per word; True  <=> non-ptr or dead
                         --                    False <=> ptr

-------------------------
-- An ArgDescr describes the argument pattern of a function

data ArgDescr
  = ArgSpec             -- Fits one of the standard patterns
        !Int            -- RTS type identifier ARG_P, ARG_N, ...

  | ArgGen              -- General case
        Liveness        -- Details about the arguments

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

mkHeapRep :: DynFlags -> IsStatic -> WordOff -> WordOff -> ClosureTypeInfo
          -> SMRep
mkHeapRep dflags is_static ptr_wds nonptr_wds cl_type_info
  = HeapRep is_static
            ptr_wds
            (nonptr_wds + slop_wds)
            cl_type_info
  where
     slop_wds
      | is_static = 0
      | otherwise = max 0 (minClosureSize dflags - (hdr_size + payload_size))

     hdr_size     = closureTypeHdrSize dflags cl_type_info
     payload_size = ptr_wds + nonptr_wds

mkRTSRep :: Int -> SMRep -> SMRep
mkRTSRep = RTSRep

mkStackRep :: [Bool] -> SMRep
mkStackRep liveness = StackRep liveness

blackHoleRep :: SMRep
blackHoleRep = HeapRep False 0 0 BlackHole

indStaticRep :: SMRep
indStaticRep = HeapRep True 1 0 IndStatic

arrPtrsRep :: DynFlags -> WordOff -> SMRep
arrPtrsRep dflags elems = ArrayPtrsRep elems (cardTableSizeW dflags elems)

smallArrPtrsRep :: WordOff -> SMRep
smallArrPtrsRep elems = SmallArrayPtrsRep elems

arrWordsRep :: DynFlags -> ByteOff -> SMRep
arrWordsRep dflags bytes = ArrayWordsRep (bytesToWordsRoundUp dflags bytes)

-----------------------------------------------------------------------------
-- Predicates

isStaticRep :: SMRep -> IsStatic
isStaticRep (HeapRep is_static _ _ _) = is_static
isStaticRep (RTSRep _ rep)            = isStaticRep rep
isStaticRep _                         = False

isStackRep :: SMRep -> Bool
isStackRep StackRep{}     = True
isStackRep (RTSRep _ rep) = isStackRep rep
isStackRep _              = False

isConRep :: SMRep -> Bool
isConRep (HeapRep _ _ _ Constr{}) = True
isConRep _                        = False

isThunkRep :: SMRep -> Bool
isThunkRep (HeapRep _ _ _ Thunk)           = True
isThunkRep (HeapRep _ _ _ ThunkSelector{}) = True
isThunkRep (HeapRep _ _ _ BlackHole)       = True
isThunkRep (HeapRep _ _ _ IndStatic)       = True
isThunkRep _                               = False

isFunRep :: SMRep -> Bool
isFunRep (HeapRep _ _ _ Fun{}) = True
isFunRep _                     = False

isStaticNoCafCon :: SMRep -> Bool
-- This should line up exactly with CONSTR_NOCAF below
-- See Note [Static NoCaf constructors]
isStaticNoCafCon (HeapRep _ 0 _ Constr{}) = True
isStaticNoCafCon _                        = False

-----------------------------------------------------------------------------
-- Size-related things

fixedHdrSize :: DynFlags -> ByteOff
fixedHdrSize dflags = wordsToBytes dflags (fixedHdrSizeW dflags)

-- | Size of a closure header (StgHeader in includes/rts/storage/Closures.h)
fixedHdrSizeW :: DynFlags -> WordOff
fixedHdrSizeW dflags = sTD_HDR_SIZE dflags + profHdrSize dflags

-- | Size of the profiling part of a closure header
-- (StgProfHeader in includes/rts/storage/Closures.h)
profHdrSize  :: DynFlags -> WordOff
profHdrSize dflags
 | gopt Opt_SccProfilingOn dflags = pROF_HDR_SIZE dflags
 | otherwise                      = 0

-- | The garbage collector requires that every closure is at least as
--   big as this.
minClosureSize :: DynFlags -> WordOff
minClosureSize dflags = fixedHdrSizeW dflags + mIN_PAYLOAD_SIZE dflags

arrWordsHdrSize :: DynFlags -> ByteOff
arrWordsHdrSize dflags
 = fixedHdrSize dflags + sIZEOF_StgArrBytes_NoHdr dflags

arrWordsHdrSizeW :: DynFlags -> WordOff
arrWordsHdrSizeW dflags =
    fixedHdrSizeW dflags +
    (sIZEOF_StgArrBytes_NoHdr dflags `quot` wORD_SIZE dflags)

arrPtrsHdrSize :: DynFlags -> ByteOff
arrPtrsHdrSize dflags
 = fixedHdrSize dflags + sIZEOF_StgMutArrPtrs_NoHdr dflags

arrPtrsHdrSizeW :: DynFlags -> WordOff
arrPtrsHdrSizeW dflags =
    fixedHdrSizeW dflags +
    (sIZEOF_StgMutArrPtrs_NoHdr dflags `quot` wORD_SIZE dflags)

smallArrPtrsHdrSize :: DynFlags -> ByteOff
smallArrPtrsHdrSize dflags
 = fixedHdrSize dflags + sIZEOF_StgSmallMutArrPtrs_NoHdr dflags

smallArrPtrsHdrSizeW :: DynFlags -> WordOff
smallArrPtrsHdrSizeW dflags =
    fixedHdrSizeW dflags +
    (sIZEOF_StgSmallMutArrPtrs_NoHdr dflags `quot` wORD_SIZE dflags)

-- Thunks have an extra header word on SMP, so the update doesn't
-- splat the payload.
thunkHdrSize :: DynFlags -> WordOff
thunkHdrSize dflags = fixedHdrSizeW dflags + smp_hdr
        where smp_hdr = sIZEOF_StgSMPThunkHeader dflags `quot` wORD_SIZE dflags

hdrSize :: DynFlags -> SMRep -> ByteOff
hdrSize dflags rep = wordsToBytes dflags (hdrSizeW dflags rep)

hdrSizeW :: DynFlags -> SMRep -> WordOff
hdrSizeW dflags (HeapRep _ _ _ ty)    = closureTypeHdrSize dflags ty
hdrSizeW dflags (ArrayPtrsRep _ _)    = arrPtrsHdrSizeW dflags
hdrSizeW dflags (SmallArrayPtrsRep _) = smallArrPtrsHdrSizeW dflags
hdrSizeW dflags (ArrayWordsRep _)     = arrWordsHdrSizeW dflags
hdrSizeW _ _                          = panic "SMRep.hdrSizeW"

nonHdrSize :: DynFlags -> SMRep -> ByteOff
nonHdrSize dflags rep = wordsToBytes dflags (nonHdrSizeW rep)

nonHdrSizeW :: SMRep -> WordOff
nonHdrSizeW (HeapRep _ p np _) = p + np
nonHdrSizeW (ArrayPtrsRep elems ct) = elems + ct
nonHdrSizeW (SmallArrayPtrsRep elems) = elems
nonHdrSizeW (ArrayWordsRep words) = words
nonHdrSizeW (StackRep bs)      = length bs
nonHdrSizeW (RTSRep _ rep)     = nonHdrSizeW rep

-- | The total size of the closure, in words.
heapClosureSizeW :: DynFlags -> SMRep -> WordOff
heapClosureSizeW dflags (HeapRep _ p np ty)
 = closureTypeHdrSize dflags ty + p + np
heapClosureSizeW dflags (ArrayPtrsRep elems ct)
 = arrPtrsHdrSizeW dflags + elems + ct
heapClosureSizeW dflags (SmallArrayPtrsRep elems)
 = smallArrPtrsHdrSizeW dflags + elems
heapClosureSizeW dflags (ArrayWordsRep words)
 = arrWordsHdrSizeW dflags + words
heapClosureSizeW _ _ = panic "SMRep.heapClosureSize"

closureTypeHdrSize :: DynFlags -> ClosureTypeInfo -> WordOff
closureTypeHdrSize dflags ty = case ty of
                  Thunk           -> thunkHdrSize dflags
                  ThunkSelector{} -> thunkHdrSize dflags
                  BlackHole       -> thunkHdrSize dflags
                  IndStatic       -> thunkHdrSize dflags
                  _               -> fixedHdrSizeW dflags
        -- All thunks use thunkHdrSize, even if they are non-updatable.
        -- this is because we don't have separate closure types for
        -- updatable vs. non-updatable thunks, so the GC can't tell the
        -- difference.  If we ever have significant numbers of non-
        -- updatable thunks, it might be worth fixing this.

-- ---------------------------------------------------------------------------
-- Arrays

-- | The byte offset into the card table of the card for a given element
card :: DynFlags -> Int -> Int
card dflags i = i `shiftR` mUT_ARR_PTRS_CARD_BITS dflags

-- | Convert a number of elements to a number of cards, rounding up
cardRoundUp :: DynFlags -> Int -> Int
cardRoundUp dflags i =
  card dflags (i + ((1 `shiftL` mUT_ARR_PTRS_CARD_BITS dflags) - 1))

-- | The size of a card table, in bytes
cardTableSizeB :: DynFlags -> Int -> ByteOff
cardTableSizeB dflags elems = cardRoundUp dflags elems

-- | The size of a card table, in words
cardTableSizeW :: DynFlags -> Int -> WordOff
cardTableSizeW dflags elems =
  bytesToWordsRoundUp dflags (cardTableSizeB dflags elems)

-----------------------------------------------------------------------------
-- deriving the RTS closure type from an SMRep

#include "../includes/rts/storage/ClosureTypes.h"
#include "../includes/rts/storage/FunTypes.h"
-- Defines CONSTR, CONSTR_1_0 etc

-- | Derives the RTS closure type from an 'SMRep'
rtsClosureType :: SMRep -> Int
rtsClosureType rep
    = case rep of
      RTSRep ty _ -> ty

      -- See Note [static constructors]
      HeapRep _     1 0 Constr{} -> CONSTR_1_0
      HeapRep _     0 1 Constr{} -> CONSTR_0_1
      HeapRep _     2 0 Constr{} -> CONSTR_2_0
      HeapRep _     1 1 Constr{} -> CONSTR_1_1
      HeapRep _     0 2 Constr{} -> CONSTR_0_2
      HeapRep _     0 _ Constr{} -> CONSTR_NOCAF
           -- See Note [Static NoCaf constructors]
      HeapRep _     _ _ Constr{} -> CONSTR

      HeapRep False 1 0 Fun{} -> FUN_1_0
      HeapRep False 0 1 Fun{} -> FUN_0_1
      HeapRep False 2 0 Fun{} -> FUN_2_0
      HeapRep False 1 1 Fun{} -> FUN_1_1
      HeapRep False 0 2 Fun{} -> FUN_0_2
      HeapRep False _ _ Fun{} -> FUN

      HeapRep False 1 0 Thunk -> THUNK_1_0
      HeapRep False 0 1 Thunk -> THUNK_0_1
      HeapRep False 2 0 Thunk -> THUNK_2_0
      HeapRep False 1 1 Thunk -> THUNK_1_1
      HeapRep False 0 2 Thunk -> THUNK_0_2
      HeapRep False _ _ Thunk -> THUNK

      HeapRep False _ _ ThunkSelector{} ->  THUNK_SELECTOR

      HeapRep True _ _ Fun{}      -> FUN_STATIC
      HeapRep True _ _ Thunk      -> THUNK_STATIC
      HeapRep False _ _ BlackHole -> BLACKHOLE
      HeapRep False _ _ IndStatic -> IND_STATIC

      _ -> panic "rtsClosureType"

-- We export these ones
rET_SMALL, rET_BIG, aRG_GEN, aRG_GEN_BIG :: Int
rET_SMALL   = RET_SMALL
rET_BIG     = RET_BIG
aRG_GEN     = ARG_GEN
aRG_GEN_BIG = ARG_GEN_BIG

{-
Note [static constructors]
~~~~~~~~~~~~~~~~~~~~~~~~~~

We used to have a CONSTR_STATIC closure type, and each constructor had
two info tables: one with CONSTR (or CONSTR_1_0 etc.), and one with
CONSTR_STATIC.

This distinction was removed, because when copying a data structure
into a compact region, we must copy static constructors into the
compact region too.  If we didn't do this, we would need to track the
references from the compact region out to the static constructors,
because they might (indirectly) refer to CAFs.

Since static constructors will be copied to the heap, if we wanted to
use different info tables for static and dynamic constructors, we
would have to switch the info pointer when copying the constructor
into the compact region, which means we would need an extra field of
the static info table to point to the dynamic one.

However, since the distinction between static and dynamic closure
types is never actually needed (other than for assertions), we can
just drop the distinction and use the same info table for both.

The GC *does* need to distinguish between static and dynamic closures,
but it does this using the HEAP_ALLOCED() macro which checks whether
the address of the closure resides within the dynamic heap.
HEAP_ALLOCED() doesn't read the closure's info table.

Note [Static NoCaf constructors]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If we know that a top-level binding 'x' is not Caffy (ie no CAFs are
reachable from 'x'), then a statically allocated constructor (Just x)
is also not Caffy, and the garbage collector need not follow its
argument fields.  Exploiting this would require two static info tables
for Just, for the two cases where the argument was Caffy or non-Caffy.

Currently we don't do this; instead we treat nullary constructors
as non-Caffy, and the others as potentially Caffy.

************************************************************************
*                                                                      *
             Pretty printing of SMRep and friends
*                                                                      *
************************************************************************
-}

instance Outputable ClosureTypeInfo where
   ppr = pprTypeInfo

instance Outputable SMRep where
   ppr (HeapRep static ps nps tyinfo)
     = hang (header <+> lbrace) 2 (ppr tyinfo <+> rbrace)
     where
       header = text "HeapRep"
                <+> if static then text "static" else empty
                <+> pp_n "ptrs" ps <+> pp_n "nonptrs" nps
       pp_n :: String -> Int -> SDoc
       pp_n _ 0 = empty
       pp_n s n = int n <+> text s

   ppr (ArrayPtrsRep size _) = text "ArrayPtrsRep" <+> ppr size

   ppr (SmallArrayPtrsRep size) = text "SmallArrayPtrsRep" <+> ppr size

   ppr (ArrayWordsRep words) = text "ArrayWordsRep" <+> ppr words

   ppr (StackRep bs) = text "StackRep" <+> ppr bs

   ppr (RTSRep ty rep) = text "tag:" <> ppr ty <+> ppr rep

instance Outputable ArgDescr where
  ppr (ArgSpec n) = text "ArgSpec" <+> ppr n
  ppr (ArgGen ls) = text "ArgGen" <+> ppr ls

pprTypeInfo :: ClosureTypeInfo -> SDoc
pprTypeInfo (Constr tag descr)
  = text "Con" <+>
    braces (sep [ text "tag:" <+> ppr tag
                , text "descr:" <> text (show descr) ])

pprTypeInfo (Fun arity args)
  = text "Fun" <+>
    braces (sep [ text "arity:" <+> ppr arity
                , ptext (sLit ("fun_type:")) <+> ppr args ])

pprTypeInfo (ThunkSelector offset)
  = text "ThunkSel" <+> ppr offset

pprTypeInfo Thunk     = text "Thunk"
pprTypeInfo BlackHole = text "BlackHole"
pprTypeInfo IndStatic = text "IndStatic"


{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Hoopl.Block
    ( C
    , O
    , MaybeO(..)
    , IndexedCO
    , Block(..)
    , blockAppend
    , blockCons
    , blockFromList
    , blockJoin
    , blockJoinHead
    , blockJoinTail
    , blockSnoc
    , blockSplit
    , blockSplitHead
    , blockSplitTail
    , blockToList
    , emptyBlock
    , firstNode
    , foldBlockNodesB
    , foldBlockNodesB3
    , foldBlockNodesF
    , isEmptyBlock
    , lastNode
    , mapBlock
    , mapBlock'
    , mapBlock3'
    , replaceFirstNode
    , replaceLastNode
    ) where

import GhcPrelude

-- -----------------------------------------------------------------------------
-- Shapes: Open and Closed

-- | Used at the type level to indicate an "open" structure with
-- a unique, unnamed control-flow edge flowing in or out.
-- "Fallthrough" and concatenation are permitted at an open point.
data O

-- | Used at the type level to indicate a "closed" structure which
-- supports control transfer only through the use of named
-- labels---no "fallthrough" is permitted.  The number of control-flow
-- edges is unconstrained.
data C

-- | Either type indexed by closed/open using type families
type family IndexedCO ex a b :: *
type instance IndexedCO C a _b = a
type instance IndexedCO O _a b = b

-- | Maybe type indexed by open/closed
data MaybeO ex t where
  JustO    :: t -> MaybeO O t
  NothingO ::      MaybeO C t

-- | Maybe type indexed by closed/open
data MaybeC ex t where
  JustC    :: t -> MaybeC C t
  NothingC ::      MaybeC O t

instance Functor (MaybeO ex) where
  fmap _ NothingO = NothingO
  fmap f (JustO a) = JustO (f a)

instance Functor (MaybeC ex) where
  fmap _ NothingC = NothingC
  fmap f (JustC a) = JustC (f a)

-- -----------------------------------------------------------------------------
-- The Block type

-- | A sequence of nodes.  May be any of four shapes (O/O, O/C, C/O, C/C).
-- Open at the entry means single entry, mutatis mutandis for exit.
-- A closed/closed block is a /basic/ block and can't be extended further.
-- Clients should avoid manipulating blocks and should stick to either nodes
-- or graphs.
data Block n e x where
  BlockCO  :: n C O -> Block n O O          -> Block n C O
  BlockCC  :: n C O -> Block n O O -> n O C -> Block n C C
  BlockOC  ::          Block n O O -> n O C -> Block n O C

  BNil    :: Block n O O
  BMiddle :: n O O                      -> Block n O O
  BCat    :: Block n O O -> Block n O O -> Block n O O
  BSnoc   :: Block n O O -> n O O       -> Block n O O
  BCons   :: n O O       -> Block n O O -> Block n O O

-- -----------------------------------------------------------------------------
-- Simple operations on Blocks

-- Predicates

isEmptyBlock :: Block n e x -> Bool
isEmptyBlock BNil       = True
isEmptyBlock (BCat l r) = isEmptyBlock l && isEmptyBlock r
isEmptyBlock _          = False

-- Building

emptyBlock :: Block n O O
emptyBlock = BNil

blockCons :: n O O -> Block n O x -> Block n O x
blockCons n b = case b of
  BlockOC b l  -> (BlockOC $! (n `blockCons` b)) l
  BNil{}    -> BMiddle n
  BMiddle{} -> n `BCons` b
  BCat{}    -> n `BCons` b
  BSnoc{}   -> n `BCons` b
  BCons{}   -> n `BCons` b

blockSnoc :: Block n e O -> n O O -> Block n e O
blockSnoc b n = case b of
  BlockCO f b -> BlockCO f $! (b `blockSnoc` n)
  BNil{}      -> BMiddle n
  BMiddle{}   -> b `BSnoc` n
  BCat{}      -> b `BSnoc` n
  BSnoc{}     -> b `BSnoc` n
  BCons{}     -> b `BSnoc` n

blockJoinHead :: n C O -> Block n O x -> Block n C x
blockJoinHead f (BlockOC b l) = BlockCC f b l
blockJoinHead f b = BlockCO f BNil `cat` b

blockJoinTail :: Block n e O -> n O C -> Block n e C
blockJoinTail (BlockCO f b) t = BlockCC f b t
blockJoinTail b t = b `cat` BlockOC BNil t

blockJoin :: n C O -> Block n O O -> n O C -> Block n C C
blockJoin f b t = BlockCC f b t

blockAppend :: Block n e O -> Block n O x -> Block n e x
blockAppend = cat

-- Taking apart

firstNode :: Block n C x -> n C O
firstNode (BlockCO n _)   = n
firstNode (BlockCC n _ _) = n

lastNode :: Block n x C -> n O C
lastNode (BlockOC   _ n) = n
lastNode (BlockCC _ _ n) = n

blockSplitHead :: Block n C x -> (n C O, Block n O x)
blockSplitHead (BlockCO n b)   = (n, b)
blockSplitHead (BlockCC n b t) = (n, BlockOC b t)

blockSplitTail :: Block n e C -> (Block n e O, n O C)
blockSplitTail (BlockOC b n)   = (b, n)
blockSplitTail (BlockCC f b t) = (BlockCO f b, t)

-- | Split a closed block into its entry node, open middle block, and
-- exit node.
blockSplit :: Block n C C -> (n C O, Block n O O, n O C)
blockSplit (BlockCC f b t) = (f, b, t)

blockToList :: Block n O O -> [n O O]
blockToList b = go b []
   where go :: Block n O O -> [n O O] -> [n O O]
         go BNil         r = r
         go (BMiddle n)  r = n : r
         go (BCat b1 b2) r = go b1 $! go b2 r
         go (BSnoc b1 n) r = go b1 (n:r)
         go (BCons n b1) r = n : go b1 r

blockFromList :: [n O O] -> Block n O O
blockFromList = foldr BCons BNil

-- Modifying

replaceFirstNode :: Block n C x -> n C O -> Block n C x
replaceFirstNode (BlockCO _ b)   f = BlockCO f b
replaceFirstNode (BlockCC _ b n) f = BlockCC f b n

replaceLastNode :: Block n x C -> n O C -> Block n x C
replaceLastNode (BlockOC   b _) n = BlockOC b n
replaceLastNode (BlockCC l b _) n = BlockCC l b n

-- -----------------------------------------------------------------------------
-- General concatenation

cat :: Block n e O -> Block n O x -> Block n e x
cat x y = case x of
  BNil -> y

  BlockCO l b1 -> case y of
                   BlockOC b2 n -> (BlockCC l $! (b1 `cat` b2)) n
                   BNil         -> x
                   BMiddle _    -> BlockCO l $! (b1 `cat` y)
                   BCat{}       -> BlockCO l $! (b1 `cat` y)
                   BSnoc{}      -> BlockCO l $! (b1 `cat` y)
                   BCons{}      -> BlockCO l $! (b1 `cat` y)

  BMiddle n -> case y of
                   BlockOC b2 n2 -> (BlockOC $! (x `cat` b2)) n2
                   BNil          -> x
                   BMiddle{}     -> BCons n y
                   BCat{}        -> BCons n y
                   BSnoc{}       -> BCons n y
                   BCons{}       -> BCons n y

  BCat{} -> case y of
                   BlockOC b3 n2 -> (BlockOC $! (x `cat` b3)) n2
                   BNil          -> x
                   BMiddle n     -> BSnoc x n
                   BCat{}        -> BCat x y
                   BSnoc{}       -> BCat x y
                   BCons{}       -> BCat x y

  BSnoc{} -> case y of
                   BlockOC b2 n2 -> (BlockOC $! (x `cat` b2)) n2
                   BNil          -> x
                   BMiddle n     -> BSnoc x n
                   BCat{}        -> BCat x y
                   BSnoc{}       -> BCat x y
                   BCons{}       -> BCat x y

  BCons{} -> case y of
                   BlockOC b2 n2 -> (BlockOC $! (x `cat` b2)) n2
                   BNil          -> x
                   BMiddle n     -> BSnoc x n
                   BCat{}        -> BCat x y
                   BSnoc{}       -> BCat x y
                   BCons{}       -> BCat x y

-- -----------------------------------------------------------------------------
-- Mapping

-- | map a function over the nodes of a 'Block'
mapBlock :: (forall e x. n e x -> n' e x) -> Block n e x -> Block n' e x
mapBlock f (BlockCO n b  ) = BlockCO (f n) (mapBlock f b)
mapBlock f (BlockOC   b n) = BlockOC       (mapBlock f b) (f n)
mapBlock f (BlockCC n b m) = BlockCC (f n) (mapBlock f b) (f m)
mapBlock _  BNil           = BNil
mapBlock f (BMiddle n)     = BMiddle (f n)
mapBlock f (BCat b1 b2)    = BCat    (mapBlock f b1) (mapBlock f b2)
mapBlock f (BSnoc b n)     = BSnoc   (mapBlock f b)  (f n)
mapBlock f (BCons n b)     = BCons   (f n)  (mapBlock f b)

-- | A strict 'mapBlock'
mapBlock' :: (forall e x. n e x -> n' e x) -> (Block n e x -> Block n' e x)
mapBlock' f = mapBlock3' (f, f, f)

-- | map over a block, with different functions to apply to first nodes,
-- middle nodes and last nodes respectively.  The map is strict.
--
mapBlock3' :: forall n n' e x .
             ( n C O -> n' C O
             , n O O -> n' O O,
               n O C -> n' O C)
          -> Block n e x -> Block n' e x
mapBlock3' (f, m, l) b = go b
  where go :: forall e x . Block n e x -> Block n' e x
        go (BlockOC b y)   = (BlockOC $! go b) $! l y
        go (BlockCO x b)   = (BlockCO $! f x) $! (go b)
        go (BlockCC x b y) = ((BlockCC $! f x) $! go b) $! (l y)
        go BNil            = BNil
        go (BMiddle n)     = BMiddle $! m n
        go (BCat x y)      = (BCat $! go x) $! (go y)
        go (BSnoc x n)     = (BSnoc $! go x) $! (m n)
        go (BCons n x)     = (BCons $! m n) $! (go x)

-- -----------------------------------------------------------------------------
-- Folding

-- | Fold a function over every node in a block, forward or backward.
-- The fold function must be polymorphic in the shape of the nodes.
foldBlockNodesF3 :: forall n a b c .
                   ( n C O       -> a -> b
                   , n O O       -> b -> b
                   , n O C       -> b -> c)
                 -> (forall e x . Block n e x -> IndexedCO e a b -> IndexedCO x c b)
foldBlockNodesF  :: forall n a .
                    (forall e x . n e x       -> a -> a)
                 -> (forall e x . Block n e x -> IndexedCO e a a -> IndexedCO x a a)
foldBlockNodesB3 :: forall n a b c .
                   ( n C O       -> b -> c
                   , n O O       -> b -> b
                   , n O C       -> a -> b)
                 -> (forall e x . Block n e x -> IndexedCO x a b -> IndexedCO e c b)
foldBlockNodesB  :: forall n a .
                    (forall e x . n e x       -> a -> a)
                 -> (forall e x . Block n e x -> IndexedCO x a a -> IndexedCO e a a)

foldBlockNodesF3 (ff, fm, fl) = block
  where block :: forall e x . Block n e x -> IndexedCO e a b -> IndexedCO x c b
        block (BlockCO f b  )   = ff f `cat` block b
        block (BlockCC f b l)   = ff f `cat` block b `cat` fl l
        block (BlockOC   b l)   =            block b `cat` fl l
        block BNil              = id
        block (BMiddle node)    = fm node
        block (b1 `BCat`    b2) = block b1 `cat` block b2
        block (b1 `BSnoc` n)    = block b1 `cat` fm n
        block (n `BCons` b2)    = fm n `cat` block b2
        cat :: forall a b c. (a -> b) -> (b -> c) -> a -> c
        cat f f' = f' . f

foldBlockNodesF f = foldBlockNodesF3 (f, f, f)

foldBlockNodesB3 (ff, fm, fl) = block
  where block :: forall e x . Block n e x -> IndexedCO x a b -> IndexedCO e c b
        block (BlockCO f b  )   = ff f `cat` block b
        block (BlockCC f b l)   = ff f `cat` block b `cat` fl l
        block (BlockOC   b l)   =            block b `cat` fl l
        block BNil              = id
        block (BMiddle node)    = fm node
        block (b1 `BCat`    b2) = block b1 `cat` block b2
        block (b1 `BSnoc` n)    = block b1 `cat` fm n
        block (n `BCons` b2)    = fm n `cat` block b2
        cat :: forall a b c. (b -> c) -> (a -> b) -> a -> c
        cat f f' = f . f'

foldBlockNodesB f = foldBlockNodesB3 (f, f, f)


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Hoopl.Collections
    ( IsSet(..)
    , setInsertList, setDeleteList, setUnions
    , IsMap(..)
    , mapInsertList, mapDeleteList, mapUnions
    , UniqueMap, UniqueSet
    ) where

import GhcPrelude

import qualified Data.IntMap.Strict as M
import qualified Data.IntSet as S

import Data.List (foldl1')

class IsSet set where
  type ElemOf set

  setNull :: set -> Bool
  setSize :: set -> Int
  setMember :: ElemOf set -> set -> Bool

  setEmpty :: set
  setSingleton :: ElemOf set -> set
  setInsert :: ElemOf set -> set -> set
  setDelete :: ElemOf set -> set -> set

  setUnion :: set -> set -> set
  setDifference :: set -> set -> set
  setIntersection :: set -> set -> set
  setIsSubsetOf :: set -> set -> Bool
  setFilter :: (ElemOf set -> Bool) -> set -> set

  setFoldl :: (b -> ElemOf set -> b) -> b -> set -> b
  setFoldr :: (ElemOf set -> b -> b) -> b -> set -> b

  setElems :: set -> [ElemOf set]
  setFromList :: [ElemOf set] -> set

-- Helper functions for IsSet class
setInsertList :: IsSet set => [ElemOf set] -> set -> set
setInsertList keys set = foldl' (flip setInsert) set keys

setDeleteList :: IsSet set => [ElemOf set] -> set -> set
setDeleteList keys set = foldl' (flip setDelete) set keys

setUnions :: IsSet set => [set] -> set
setUnions [] = setEmpty
setUnions sets = foldl1' setUnion sets

class IsMap map where
  type KeyOf map

  mapNull :: map a -> Bool
  mapSize :: map a -> Int
  mapMember :: KeyOf map -> map a -> Bool
  mapLookup :: KeyOf map -> map a -> Maybe a
  mapFindWithDefault :: a -> KeyOf map -> map a -> a

  mapEmpty :: map a
  mapSingleton :: KeyOf map -> a -> map a
  mapInsert :: KeyOf map -> a -> map a -> map a
  mapInsertWith :: (a -> a -> a) -> KeyOf map -> a -> map a -> map a
  mapDelete :: KeyOf map -> map a -> map a
  mapAlter :: (Maybe a -> Maybe a) -> KeyOf map -> map a -> map a
  mapAdjust :: (a -> a) -> KeyOf map -> map a -> map a

  mapUnion :: map a -> map a -> map a
  mapUnionWithKey :: (KeyOf map -> a -> a -> a) -> map a -> map a -> map a
  mapDifference :: map a -> map a -> map a
  mapIntersection :: map a -> map a -> map a
  mapIsSubmapOf :: Eq a => map a -> map a -> Bool

  mapMap :: (a -> b) -> map a -> map b
  mapMapWithKey :: (KeyOf map -> a -> b) -> map a -> map b
  mapFoldl :: (b -> a -> b) -> b -> map a -> b
  mapFoldr :: (a -> b -> b) -> b -> map a -> b
  mapFoldlWithKey :: (b -> KeyOf map -> a -> b) -> b -> map a -> b
  mapFoldMapWithKey :: Monoid m => (KeyOf map -> a -> m) -> map a -> m
  mapFilter :: (a -> Bool) -> map a -> map a
  mapFilterWithKey :: (KeyOf map -> a -> Bool) -> map a -> map a

  mapElems :: map a -> [a]
  mapKeys :: map a -> [KeyOf map]
  mapToList :: map a -> [(KeyOf map, a)]
  mapFromList :: [(KeyOf map, a)] -> map a
  mapFromListWith :: (a -> a -> a) -> [(KeyOf map,a)] -> map a

-- Helper functions for IsMap class
mapInsertList :: IsMap map => [(KeyOf map, a)] -> map a -> map a
mapInsertList assocs map = foldl' (flip (uncurry mapInsert)) map assocs

mapDeleteList :: IsMap map => [KeyOf map] -> map a -> map a
mapDeleteList keys map = foldl' (flip mapDelete) map keys

mapUnions :: IsMap map => [map a] -> map a
mapUnions [] = mapEmpty
mapUnions maps = foldl1' mapUnion maps

-----------------------------------------------------------------------------
-- Basic instances
-----------------------------------------------------------------------------

newtype UniqueSet = US S.IntSet deriving (Eq, Ord, Show, Semigroup, Monoid)

instance IsSet UniqueSet where
  type ElemOf UniqueSet = Int

  setNull (US s) = S.null s
  setSize (US s) = S.size s
  setMember k (US s) = S.member k s

  setEmpty = US S.empty
  setSingleton k = US (S.singleton k)
  setInsert k (US s) = US (S.insert k s)
  setDelete k (US s) = US (S.delete k s)

  setUnion (US x) (US y) = US (S.union x y)
  setDifference (US x) (US y) = US (S.difference x y)
  setIntersection (US x) (US y) = US (S.intersection x y)
  setIsSubsetOf (US x) (US y) = S.isSubsetOf x y
  setFilter f (US s) = US (S.filter f s)

  setFoldl k z (US s) = S.foldl' k z s
  setFoldr k z (US s) = S.foldr k z s

  setElems (US s) = S.elems s
  setFromList ks = US (S.fromList ks)

newtype UniqueMap v = UM (M.IntMap v)
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

instance IsMap UniqueMap where
  type KeyOf UniqueMap = Int

  mapNull (UM m) = M.null m
  mapSize (UM m) = M.size m
  mapMember k (UM m) = M.member k m
  mapLookup k (UM m) = M.lookup k m
  mapFindWithDefault def k (UM m) = M.findWithDefault def k m

  mapEmpty = UM M.empty
  mapSingleton k v = UM (M.singleton k v)
  mapInsert k v (UM m) = UM (M.insert k v m)
  mapInsertWith f k v (UM m) = UM (M.insertWith f k v m)
  mapDelete k (UM m) = UM (M.delete k m)
  mapAlter f k (UM m) = UM (M.alter f k m)
  mapAdjust f k (UM m) = UM (M.adjust f k m)

  mapUnion (UM x) (UM y) = UM (M.union x y)
  mapUnionWithKey f (UM x) (UM y) = UM (M.unionWithKey f x y)
  mapDifference (UM x) (UM y) = UM (M.difference x y)
  mapIntersection (UM x) (UM y) = UM (M.intersection x y)
  mapIsSubmapOf (UM x) (UM y) = M.isSubmapOf x y

  mapMap f (UM m) = UM (M.map f m)
  mapMapWithKey f (UM m) = UM (M.mapWithKey f m)
  mapFoldl k z (UM m) = M.foldl' k z m
  mapFoldr k z (UM m) = M.foldr k z m
  mapFoldlWithKey k z (UM m) = M.foldlWithKey' k z m
  mapFoldMapWithKey f (UM m) = M.foldMapWithKey f m
  mapFilter f (UM m) = UM (M.filter f m)
  mapFilterWithKey f (UM m) = UM (M.filterWithKey f m)

  mapElems (UM m) = M.elems m
  mapKeys (UM m) = M.keys m
  mapToList (UM m) = M.toList m
  mapFromList assocs = UM (M.fromList assocs)
  mapFromListWith f assocs = UM (M.fromListWith f assocs)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fprof-auto-top #-}

--
-- Copyright (c) 2010, JoÃ£o Dias, Simon Marlow, Simon Peyton Jones,
-- and Norman Ramsey
--
-- Modifications copyright (c) The University of Glasgow 2012
--
-- This module is a specialised and optimised version of
-- Compiler.Hoopl.Dataflow in the hoopl package.  In particular it is
-- specialised to the UniqSM monad.
--

module Hoopl.Dataflow
  ( C, O, Block
  , lastNode, entryLabel
  , foldNodesBwdOO
  , foldRewriteNodesBwdOO
  , DataflowLattice(..), OldFact(..), NewFact(..), JoinedFact(..)
  , TransferFun, RewriteFun
  , Fact, FactBase
  , getFact, mkFactBase
  , analyzeCmmFwd, analyzeCmmBwd
  , rewriteCmmBwd
  , changedIf
  , joinOutFacts
  , joinFacts
  )
where

import GhcPrelude

import Cmm
import UniqSupply

import Data.Array
import Data.Maybe
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet

import Hoopl.Block
import Hoopl.Graph
import Hoopl.Collections
import Hoopl.Label

type family   Fact x f :: *
type instance Fact C f = FactBase f
type instance Fact O f = f

newtype OldFact a = OldFact a

newtype NewFact a = NewFact a

-- | The result of joining OldFact and NewFact.
data JoinedFact a
    = Changed !a     -- ^ Result is different than OldFact.
    | NotChanged !a  -- ^ Result is the same as OldFact.

getJoined :: JoinedFact a -> a
getJoined (Changed a) = a
getJoined (NotChanged a) = a

changedIf :: Bool -> a -> JoinedFact a
changedIf True = Changed
changedIf False = NotChanged

type JoinFun a = OldFact a -> NewFact a -> JoinedFact a

data DataflowLattice a = DataflowLattice
    { fact_bot :: a
    , fact_join :: JoinFun a
    }

data Direction = Fwd | Bwd

type TransferFun f = CmmBlock -> FactBase f -> FactBase f

-- | Function for rewrtiting and analysis combined. To be used with
-- @rewriteCmm@.
--
-- Currently set to work with @UniqSM@ monad, but we could probably abstract
-- that away (if we do that, we might want to specialize the fixpoint algorithms
-- to the particular monads through SPECIALIZE).
type RewriteFun f = CmmBlock -> FactBase f -> UniqSM (CmmBlock, FactBase f)

analyzeCmmBwd, analyzeCmmFwd
    :: DataflowLattice f
    -> TransferFun f
    -> CmmGraph
    -> FactBase f
    -> FactBase f
analyzeCmmBwd = analyzeCmm Bwd
analyzeCmmFwd = analyzeCmm Fwd

analyzeCmm
    :: Direction
    -> DataflowLattice f
    -> TransferFun f
    -> CmmGraph
    -> FactBase f
    -> FactBase f
analyzeCmm dir lattice transfer cmmGraph initFact =
    let entry = g_entry cmmGraph
        hooplGraph = g_graph cmmGraph
        blockMap =
            case hooplGraph of
                GMany NothingO bm NothingO -> bm
    in fixpointAnalysis dir lattice transfer entry blockMap initFact

-- Fixpoint algorithm.
fixpointAnalysis
    :: forall f.
       Direction
    -> DataflowLattice f
    -> TransferFun f
    -> Label
    -> LabelMap CmmBlock
    -> FactBase f
    -> FactBase f
fixpointAnalysis direction lattice do_block entry blockmap = loop start
  where
    -- Sorting the blocks helps to minimize the number of times we need to
    -- process blocks. For instance, for forward analysis we want to look at
    -- blocks in reverse postorder. Also, see comments for sortBlocks.
    blocks     = sortBlocks direction entry blockmap
    num_blocks = length blocks
    block_arr  = {-# SCC "block_arr" #-} listArray (0, num_blocks - 1) blocks
    start      = {-# SCC "start" #-} IntSet.fromDistinctAscList
      [0 .. num_blocks - 1]
    dep_blocks = {-# SCC "dep_blocks" #-} mkDepBlocks direction blocks
    join       = fact_join lattice

    loop
        :: IntHeap     -- ^ Worklist, i.e., blocks to process
        -> FactBase f  -- ^ Current result (increases monotonically)
        -> FactBase f
    loop todo !fbase1 | Just (index, todo1) <- IntSet.minView todo =
        let block = block_arr ! index
            out_facts = {-# SCC "do_block" #-} do_block block fbase1
            -- For each of the outgoing edges, we join it with the current
            -- information in fbase1 and (if something changed) we update it
            -- and add the affected blocks to the worklist.
            (todo2, fbase2) = {-# SCC "mapFoldWithKey" #-}
                mapFoldlWithKey
                    (updateFact join dep_blocks) (todo1, fbase1) out_facts
        in loop todo2 fbase2
    loop _ !fbase1 = fbase1

rewriteCmmBwd
    :: DataflowLattice f
    -> RewriteFun f
    -> CmmGraph
    -> FactBase f
    -> UniqSM (CmmGraph, FactBase f)
rewriteCmmBwd = rewriteCmm Bwd

rewriteCmm
    :: Direction
    -> DataflowLattice f
    -> RewriteFun f
    -> CmmGraph
    -> FactBase f
    -> UniqSM (CmmGraph, FactBase f)
rewriteCmm dir lattice rwFun cmmGraph initFact = do
    let entry = g_entry cmmGraph
        hooplGraph = g_graph cmmGraph
        blockMap1 =
            case hooplGraph of
                GMany NothingO bm NothingO -> bm
    (blockMap2, facts) <-
        fixpointRewrite dir lattice rwFun entry blockMap1 initFact
    return (cmmGraph {g_graph = GMany NothingO blockMap2 NothingO}, facts)

fixpointRewrite
    :: forall f.
       Direction
    -> DataflowLattice f
    -> RewriteFun f
    -> Label
    -> LabelMap CmmBlock
    -> FactBase f
    -> UniqSM (LabelMap CmmBlock, FactBase f)
fixpointRewrite dir lattice do_block entry blockmap = loop start blockmap
  where
    -- Sorting the blocks helps to minimize the number of times we need to
    -- process blocks. For instance, for forward analysis we want to look at
    -- blocks in reverse postorder. Also, see comments for sortBlocks.
    blocks     = sortBlocks dir entry blockmap
    num_blocks = length blocks
    block_arr  = {-# SCC "block_arr_rewrite" #-}
                 listArray (0, num_blocks - 1) blocks
    start      = {-# SCC "start_rewrite" #-}
                 IntSet.fromDistinctAscList [0 .. num_blocks - 1]
    dep_blocks = {-# SCC "dep_blocks_rewrite" #-} mkDepBlocks dir blocks
    join       = fact_join lattice

    loop
        :: IntHeap            -- ^ Worklist, i.e., blocks to process
        -> LabelMap CmmBlock  -- ^ Rewritten blocks.
        -> FactBase f         -- ^ Current facts.
        -> UniqSM (LabelMap CmmBlock, FactBase f)
    loop todo !blocks1 !fbase1
      | Just (index, todo1) <- IntSet.minView todo = do
        -- Note that we use the *original* block here. This is important.
        -- We're optimistically rewriting blocks even before reaching the fixed
        -- point, which means that the rewrite might be incorrect. So if the
        -- facts change, we need to rewrite the original block again (taking
        -- into account the new facts).
        let block = block_arr ! index
        (new_block, out_facts) <- {-# SCC "do_block_rewrite" #-}
            do_block block fbase1
        let blocks2 = mapInsert (entryLabel new_block) new_block blocks1
            (todo2, fbase2) = {-# SCC "mapFoldWithKey_rewrite" #-}
                mapFoldlWithKey
                    (updateFact join dep_blocks) (todo1, fbase1) out_facts
        loop todo2 blocks2 fbase2
    loop _ !blocks1 !fbase1 = return (blocks1, fbase1)

{-
Note [Unreachable blocks]
~~~~~~~~~~~~~~~~~~~~~~~~~
A block that is not in the domain of tfb_fbase is "currently unreachable".
A currently-unreachable block is not even analyzed.  Reason: consider
constant prop and this graph, with entry point L1:
  L1: x:=3; goto L4
  L2: x:=4; goto L4
  L4: if x>3 goto L2 else goto L5
Here L2 is actually unreachable, but if we process it with bottom input fact,
we'll propagate (x=4) to L4, and nuke the otherwise-good rewriting of L4.

* If a currently-unreachable block is not analyzed, then its rewritten
  graph will not be accumulated in tfb_rg.  And that is good:
  unreachable blocks simply do not appear in the output.

* Note that clients must be careful to provide a fact (even if bottom)
  for each entry point. Otherwise useful blocks may be garbage collected.

* Note that updateFact must set the change-flag if a label goes from
  not-in-fbase to in-fbase, even if its fact is bottom.  In effect the
  real fact lattice is
       UNR
       bottom
       the points above bottom

* Even if the fact is going from UNR to bottom, we still call the
  client's fact_join function because it might give the client
  some useful debugging information.

* All of this only applies for *forward* ixpoints.  For the backward
  case we must treat every block as reachable; it might finish with a
  'return', and therefore have no successors, for example.
-}

-----------------------------------------------------------------------------
--  Pieces that are shared by fixpoint and fixpoint_anal
-----------------------------------------------------------------------------

-- | Sort the blocks into the right order for analysis. This means reverse
-- postorder for a forward analysis. For the backward one, we simply reverse
-- that (see Note [Backward vs forward analysis]).
sortBlocks
    :: NonLocal n
    => Direction -> Label -> LabelMap (Block n C C) -> [Block n C C]
sortBlocks direction entry blockmap =
    case direction of
        Fwd -> fwd
        Bwd -> reverse fwd
  where
    fwd = revPostorderFrom blockmap entry

-- Note [Backward vs forward analysis]
--
-- The forward and backward cases are not dual.  In the forward case, the entry
-- points are known, and one simply traverses the body blocks from those points.
-- In the backward case, something is known about the exit points, but a
-- backward analysis must also include reachable blocks that don't reach the
-- exit, as in a procedure that loops forever and has side effects.)
-- For instance, let E be the entry and X the exit blocks (arrows indicate
-- control flow)
--   E -> X
--   E -> B
--   B -> C
--   C -> B
-- We do need to include B and C even though they're unreachable in the
-- *reverse* graph (that we could use for backward analysis):
--   E <- X
--   E <- B
--   B <- C
--   C <- B
-- So when sorting the blocks for the backward analysis, we simply take the
-- reverse of what is used for the forward one.

-- | Construct a mapping from a @Label@ to the block indexes that should be
-- re-analyzed if the facts at that @Label@ change.
--
-- Note that we're considering here the entry point of the block, so if the
-- facts change at the entry:
-- * for a backward analysis we need to re-analyze all the predecessors, but
-- * for a forward analysis, we only need to re-analyze the current block
--   (and that will in turn propagate facts into its successors).
mkDepBlocks :: Direction -> [CmmBlock] -> LabelMap IntSet
mkDepBlocks Fwd blocks = go blocks 0 mapEmpty
  where
    go []     !_ !dep_map = dep_map
    go (b:bs) !n !dep_map =
        go bs (n + 1) $ mapInsert (entryLabel b) (IntSet.singleton n) dep_map
mkDepBlocks Bwd blocks = go blocks 0 mapEmpty
  where
    go []     !_ !dep_map = dep_map
    go (b:bs) !n !dep_map =
        let insert m l = mapInsertWith IntSet.union l (IntSet.singleton n) m
        in go bs (n + 1) $ foldl' insert dep_map (successors b)

-- | After some new facts have been generated by analysing a block, we
-- fold this function over them to generate (a) a list of block
-- indices to (re-)analyse, and (b) the new FactBase.
updateFact
    :: JoinFun f
    -> LabelMap IntSet
    -> (IntHeap, FactBase f)
    -> Label
    -> f -- out fact
    -> (IntHeap, FactBase f)
updateFact fact_join dep_blocks (todo, fbase) lbl new_fact
  = case lookupFact lbl fbase of
      Nothing ->
          -- Note [No old fact]
          let !z = mapInsert lbl new_fact fbase in (changed, z)
      Just old_fact ->
          case fact_join (OldFact old_fact) (NewFact new_fact) of
              (NotChanged _) -> (todo, fbase)
              (Changed f) -> let !z = mapInsert lbl f fbase in (changed, z)
  where
    changed = todo `IntSet.union`
              mapFindWithDefault IntSet.empty lbl dep_blocks

{-
Note [No old fact]

We know that the new_fact is >= _|_, so we don't need to join.  However,
if the new fact is also _|_, and we have already analysed its block,
we don't need to record a change.  So there's a tradeoff here.  It turns
out that always recording a change is faster.
-}

----------------------------------------------------------------
--       Utilities
----------------------------------------------------------------

-- Fact lookup: the fact `orelse` bottom
getFact  :: DataflowLattice f -> Label -> FactBase f -> f
getFact lat l fb = case lookupFact l fb of Just  f -> f
                                           Nothing -> fact_bot lat

-- | Returns the result of joining the facts from all the successors of the
-- provided node or block.
joinOutFacts :: (NonLocal n) => DataflowLattice f -> n e C -> FactBase f -> f
joinOutFacts lattice nonLocal fact_base = foldl' join (fact_bot lattice) facts
  where
    join new old = getJoined $ fact_join lattice (OldFact old) (NewFact new)
    facts =
        [ fromJust fact
        | s <- successors nonLocal
        , let fact = lookupFact s fact_base
        , isJust fact
        ]

joinFacts :: DataflowLattice f -> [f] -> f
joinFacts lattice facts  = foldl' join (fact_bot lattice) facts
  where
    join new old = getJoined $ fact_join lattice (OldFact old) (NewFact new)

-- | Returns the joined facts for each label.
mkFactBase :: DataflowLattice f -> [(Label, f)] -> FactBase f
mkFactBase lattice = foldl' add mapEmpty
  where
    join = fact_join lattice

    add result (l, f1) =
        let !newFact =
                case mapLookup l result of
                    Nothing -> f1
                    Just f2 -> getJoined $ join (OldFact f1) (NewFact f2)
        in mapInsert l newFact result

-- | Folds backward over all nodes of an open-open block.
-- Strict in the accumulator.
foldNodesBwdOO :: (CmmNode O O -> f -> f) -> Block CmmNode O O -> f -> f
foldNodesBwdOO funOO = go
  where
    go (BCat b1 b2) f = go b1 $! go b2 f
    go (BSnoc h n) f = go h $! funOO n f
    go (BCons n t) f = funOO n $! go t f
    go (BMiddle n) f = funOO n f
    go BNil f = f
{-# INLINABLE foldNodesBwdOO #-}

-- | Folds backward over all the nodes of an open-open block and allows
-- rewriting them. The accumulator is both the block of nodes and @f@ (usually
-- dataflow facts).
-- Strict in both accumulated parts.
foldRewriteNodesBwdOO
    :: forall f.
       (CmmNode O O -> f -> UniqSM (Block CmmNode O O, f))
    -> Block CmmNode O O
    -> f
    -> UniqSM (Block CmmNode O O, f)
foldRewriteNodesBwdOO rewriteOO initBlock initFacts = go initBlock initFacts
  where
    go (BCons node1 block1) !fact1 = (rewriteOO node1 `comp` go block1) fact1
    go (BSnoc block1 node1) !fact1 = (go block1 `comp` rewriteOO node1) fact1
    go (BCat blockA1 blockB1) !fact1 = (go blockA1 `comp` go blockB1) fact1
    go (BMiddle node) !fact1 = rewriteOO node fact1
    go BNil !fact = return (BNil, fact)

    comp rew1 rew2 = \f1 -> do
        (b, f2) <- rew2 f1
        (a, !f3) <- rew1 f2
        let !c = joinBlocksOO a b
        return (c, f3)
    {-# INLINE comp #-}
{-# INLINABLE foldRewriteNodesBwdOO #-}

joinBlocksOO :: Block n O O -> Block n O O -> Block n O O
joinBlocksOO BNil b = b
joinBlocksOO b BNil = b
joinBlocksOO (BMiddle n) b = blockCons n b
joinBlocksOO b (BMiddle n) = blockSnoc b n
joinBlocksOO b1 b2 = BCat b1 b2

type IntHeap = IntSet


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Hoopl.Graph
    ( Body
    , Graph
    , Graph'(..)
    , NonLocal(..)
    , addBlock
    , bodyList
    , emptyBody
    , labelsDefined
    , mapGraph
    , mapGraphBlocks
    , revPostorderFrom
    ) where

import GhcPrelude
import Util

import Hoopl.Label
import Hoopl.Block
import Hoopl.Collections

-- | A (possibly empty) collection of closed/closed blocks
type Body n = LabelMap (Block n C C)

-- | @Body@ abstracted over @block@
type Body' block (n :: * -> * -> *) = LabelMap (block n C C)

-------------------------------
-- | Gives access to the anchor points for
-- nonlocal edges as well as the edges themselves
class NonLocal thing where
  entryLabel :: thing C x -> Label   -- ^ The label of a first node or block
  successors :: thing e C -> [Label] -- ^ Gives control-flow successors

instance NonLocal n => NonLocal (Block n) where
  entryLabel (BlockCO f _)   = entryLabel f
  entryLabel (BlockCC f _ _) = entryLabel f

  successors (BlockOC   _ n) = successors n
  successors (BlockCC _ _ n) = successors n

emptyBody :: Body' block n
emptyBody = mapEmpty

bodyList :: Body' block n -> [(Label,block n C C)]
bodyList body = mapToList body

addBlock
    :: (NonLocal block, HasDebugCallStack)
    => block C C -> LabelMap (block C C) -> LabelMap (block C C)
addBlock block body = mapAlter add lbl body
  where
    lbl = entryLabel block
    add Nothing = Just block
    add _ = error $ "duplicate label " ++ show lbl ++ " in graph"

-- ---------------------------------------------------------------------------
-- Graph

-- | A control-flow graph, which may take any of four shapes (O/O,
-- O/C, C/O, C/C).  A graph open at the entry has a single,
-- distinguished, anonymous entry point; if a graph is closed at the
-- entry, its entry point(s) are supplied by a context.
type Graph = Graph' Block

-- | @Graph'@ is abstracted over the block type, so that we can build
-- graphs of annotated blocks for example (Compiler.Hoopl.Dataflow
-- needs this).
data Graph' block (n :: * -> * -> *) e x where
  GNil  :: Graph' block n O O
  GUnit :: block n O O -> Graph' block n O O
  GMany :: MaybeO e (block n O C)
        -> Body' block n
        -> MaybeO x (block n C O)
        -> Graph' block n e x

-- -----------------------------------------------------------------------------
-- Mapping over graphs

-- | Maps over all nodes in a graph.
mapGraph :: (forall e x. n e x -> n' e x) -> Graph n e x -> Graph n' e x
mapGraph f = mapGraphBlocks (mapBlock f)

-- | Function 'mapGraphBlocks' enables a change of representation of blocks,
-- nodes, or both.  It lifts a polymorphic block transform into a polymorphic
-- graph transform.  When the block representation stabilizes, a similar
-- function should be provided for blocks.
mapGraphBlocks :: forall block n block' n' e x .
                  (forall e x . block n e x -> block' n' e x)
               -> (Graph' block n e x -> Graph' block' n' e x)

mapGraphBlocks f = map
  where map :: Graph' block n e x -> Graph' block' n' e x
        map GNil = GNil
        map (GUnit b) = GUnit (f b)
        map (GMany e b x) = GMany (fmap f e) (mapMap f b) (fmap f x)

-- -----------------------------------------------------------------------------
-- Extracting Labels from graphs

labelsDefined :: forall block n e x . NonLocal (block n) => Graph' block n e x
              -> LabelSet
labelsDefined GNil      = setEmpty
labelsDefined (GUnit{}) = setEmpty
labelsDefined (GMany _ body x) = mapFoldlWithKey addEntry (exitLabel x) body
  where addEntry :: forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
        addEntry labels label _ = setInsert label labels
        exitLabel :: MaybeO x (block n C O) -> LabelSet
        exitLabel NothingO  = setEmpty
        exitLabel (JustO b) = setSingleton (entryLabel b)

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

-- | Returns a list of blocks reachable from the provided Labels in the reverse
-- postorder.
--
-- This is the most important traversal over this data structure.  It drops
-- unreachable code and puts blocks in an order that is good for solving forward
-- dataflow problems quickly.  The reverse order is good for solving backward
-- dataflow problems quickly.  The forward order is also reasonably good for
-- emitting instructions, except that it will not usually exploit Forrest
-- Baskett's trick of eliminating the unconditional branch from a loop.  For
-- that you would need a more serious analysis, probably based on dominators, to
-- identify loop headers.
--
-- For forward analyses we want reverse postorder visitation, consider:
-- @
--      A -> [B,C]
--      B -> D
--      C -> D
-- @
-- Postorder: [D, C, B, A] (or [D, B, C, A])
-- Reverse postorder: [A, B, C, D] (or [A, C, B, D])
-- This matters for, e.g., forward analysis, because we want to analyze *both*
-- B and C before we analyze D.
revPostorderFrom
  :: forall block.  (NonLocal block)
  => LabelMap (block C C) -> Label -> [block C C]
revPostorderFrom graph start = go start_worklist setEmpty []
  where
    start_worklist = lookup_for_descend start Nil

    -- To compute the postorder we need to "visit" a block (mark as done)
    -- *after* visiting all its successors. So we need to know whether we
    -- already processed all successors of each block (and @NonLocal@ allows
    -- arbitrary many successors). So we use an explicit stack with an extra bit
    -- of information:
    -- * @ConsTodo@ means to explore the block if it wasn't visited before
    -- * @ConsMark@ means that all successors were already done and we can add
    --   the block to the result.
    --
    -- NOTE: We add blocks to the result list in postorder, but we *prepend*
    -- them (i.e., we use @(:)@), which means that the final list is in reverse
    -- postorder.
    go :: DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
    go Nil                      !_           !result = result
    go (ConsMark block rest)    !wip_or_done !result =
        go rest wip_or_done (block : result)
    go (ConsTodo block rest)    !wip_or_done !result
        | entryLabel block `setMember` wip_or_done = go rest wip_or_done result
        | otherwise =
            let new_worklist =
                    foldr lookup_for_descend
                          (ConsMark block rest)
                          (successors block)
            in go new_worklist (setInsert (entryLabel block) wip_or_done) result

    lookup_for_descend :: Label -> DfsStack (block C C) -> DfsStack (block C C)
    lookup_for_descend label wl
      | Just b <- mapLookup label graph = ConsTodo b wl
      | otherwise =
           error $ "Label that doesn't have a block?! " ++ show label

data DfsStack a = ConsTodo a (DfsStack a) | ConsMark a (DfsStack a) | Nil


{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Hoopl.Label
    ( Label
    , LabelMap
    , LabelSet
    , FactBase
    , lookupFact
    , mkHooplLabel
    ) where

import GhcPrelude

import Outputable

-- TODO: This should really just use GHC's Unique and Uniq{Set,FM}
import Hoopl.Collections

import Unique (Uniquable(..))
import TrieMap

-----------------------------------------------------------------------------
--              Label
-----------------------------------------------------------------------------

newtype Label = Label { lblToUnique :: Int }
  deriving (Eq, Ord)

mkHooplLabel :: Int -> Label
mkHooplLabel = Label

instance Show Label where
  show (Label n) = "L" ++ show n

instance Uniquable Label where
  getUnique label = getUnique (lblToUnique label)

instance Outputable Label where
  ppr label = ppr (getUnique label)

-----------------------------------------------------------------------------
-- LabelSet

newtype LabelSet = LS UniqueSet deriving (Eq, Ord, Show, Monoid, Semigroup)

instance IsSet LabelSet where
  type ElemOf LabelSet = Label

  setNull (LS s) = setNull s
  setSize (LS s) = setSize s
  setMember (Label k) (LS s) = setMember k s

  setEmpty = LS setEmpty
  setSingleton (Label k) = LS (setSingleton k)
  setInsert (Label k) (LS s) = LS (setInsert k s)
  setDelete (Label k) (LS s) = LS (setDelete k s)

  setUnion (LS x) (LS y) = LS (setUnion x y)
  setDifference (LS x) (LS y) = LS (setDifference x y)
  setIntersection (LS x) (LS y) = LS (setIntersection x y)
  setIsSubsetOf (LS x) (LS y) = setIsSubsetOf x y
  setFilter f (LS s) = LS (setFilter (f . mkHooplLabel) s)
  setFoldl k z (LS s) = setFoldl (\a v -> k a (mkHooplLabel v)) z s
  setFoldr k z (LS s) = setFoldr (\v a -> k (mkHooplLabel v) a) z s

  setElems (LS s) = map mkHooplLabel (setElems s)
  setFromList ks = LS (setFromList (map lblToUnique ks))

-----------------------------------------------------------------------------
-- LabelMap

newtype LabelMap v = LM (UniqueMap v)
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

instance IsMap LabelMap where
  type KeyOf LabelMap = Label

  mapNull (LM m) = mapNull m
  mapSize (LM m) = mapSize m
  mapMember (Label k) (LM m) = mapMember k m
  mapLookup (Label k) (LM m) = mapLookup k m
  mapFindWithDefault def (Label k) (LM m) = mapFindWithDefault def k m

  mapEmpty = LM mapEmpty
  mapSingleton (Label k) v = LM (mapSingleton k v)
  mapInsert (Label k) v (LM m) = LM (mapInsert k v m)
  mapInsertWith f (Label k) v (LM m) = LM (mapInsertWith f k v m)
  mapDelete (Label k) (LM m) = LM (mapDelete k m)
  mapAlter f (Label k) (LM m) = LM (mapAlter f k m)
  mapAdjust f (Label k) (LM m) = LM (mapAdjust f k m)

  mapUnion (LM x) (LM y) = LM (mapUnion x y)
  mapUnionWithKey f (LM x) (LM y) = LM (mapUnionWithKey (f . mkHooplLabel) x y)
  mapDifference (LM x) (LM y) = LM (mapDifference x y)
  mapIntersection (LM x) (LM y) = LM (mapIntersection x y)
  mapIsSubmapOf (LM x) (LM y) = mapIsSubmapOf x y

  mapMap f (LM m) = LM (mapMap f m)
  mapMapWithKey f (LM m) = LM (mapMapWithKey (f . mkHooplLabel) m)
  mapFoldl k z (LM m) = mapFoldl k z m
  mapFoldr k z (LM m) = mapFoldr k z m
  mapFoldlWithKey k z (LM m) =
      mapFoldlWithKey (\a v -> k a (mkHooplLabel v)) z m
  mapFoldMapWithKey f (LM m) = mapFoldMapWithKey (\k v -> f (mkHooplLabel k) v) m
  mapFilter f (LM m) = LM (mapFilter f m)
  mapFilterWithKey f (LM m) = LM (mapFilterWithKey (f . mkHooplLabel) m)

  mapElems (LM m) = mapElems m
  mapKeys (LM m) = map mkHooplLabel (mapKeys m)
  mapToList (LM m) = [(mkHooplLabel k, v) | (k, v) <- mapToList m]
  mapFromList assocs = LM (mapFromList [(lblToUnique k, v) | (k, v) <- assocs])
  mapFromListWith f assocs = LM (mapFromListWith f [(lblToUnique k, v) | (k, v) <- assocs])

-----------------------------------------------------------------------------
-- Instances

instance Outputable LabelSet where
  ppr = ppr . setElems

instance Outputable a => Outputable (LabelMap a) where
  ppr = ppr . mapToList

instance TrieMap LabelMap where
  type Key LabelMap = Label
  emptyTM = mapEmpty
  lookupTM k m = mapLookup k m
  alterTM k f m = mapAlter f k m
  foldTM k m z = mapFoldr k z m
  mapTM f m = mapMap f m

-----------------------------------------------------------------------------
-- FactBase

type FactBase f = LabelMap f

lookupFact :: Label -> FactBase f -> Maybe f
lookupFact = mapLookup