{-# LANGUAGE LambdaCase, RecordWildCards, MagicHash, UnboxedTuples, PatternSynonyms, ExplicitNamespaces #-}
module GHC.Cmm.UniqueRenamer
  ( detRenameCmmGroup
  , detRenameIPEMap
  , MonadGetUnique(..)

  -- Careful! Not for general use!
  , DetUniqFM, emptyDetUFM

  , module GHC.Types.Unique.DSM
  )
  where

import GHC.Prelude
import GHC.Utils.Monad.State.Strict
import Data.Tuple (swap)
import GHC.Word
import GHC.Cmm
import GHC.Cmm.CLabel
import GHC.Cmm.Dataflow.Block
import GHC.Cmm.Dataflow.Graph
import GHC.Cmm.Dataflow.Label
import GHC.Cmm.Switch
import GHC.Types.Unique
import GHC.Types.Unique.FM
import GHC.Types.Unique.DFM
import GHC.Utils.Outputable as Outputable
import GHC.Types.Id
import GHC.Types.Unique.DSM
import GHC.Types.Name hiding (varName)
import GHC.Types.Var
import GHC.Types.IPE

{-
Note [Renaming uniques deterministically]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
As mentioned by Note [Object determinism], a key step in producing
deterministic objects is to rename all existing uniques deterministically.

An important observation is that GHC already produces code in a deterministic
order, both declarations (say, A_closure always comes before B_closure) and the
instructions and data within.

We can leverage this /deterministic order/ to
rename all uniques deterministically, by traversing, specifically, Cmm code
fresh off of StgToCmm and assigning a new unique from a deterministic supply
(an incrementing counter) to every non-external unique in the order they are found.

Since the order is deterministic across runs, so will the renamed uniques.

This Cmm renaming pass is guarded by -fobject-determinism because it means the
compiler must do more work. However, performance profiling has shown the impact
to be small enough that we should consider enabling -fobject-determinism by
default instead eventually.
-}

-- | A mapping from non-deterministic uniques to deterministic uniques, to
-- rename local symbols with the end goal of producing deterministic object files.
-- See Note [Renaming uniques deterministically]
data DetUniqFM = DetUniqFM
  { mapping :: !(UniqFM Unique Unique)
  , supply  :: !Word64
  }

instance Outputable DetUniqFM where
  ppr DetUniqFM{mapping, supply} =
    ppr mapping $$
    text "supply:" Outputable.<> ppr supply

type DetRnM = State DetUniqFM

emptyDetUFM :: DetUniqFM
emptyDetUFM = DetUniqFM
  { mapping = emptyUFM
  -- NB: A lower initial value can get us label `Lsl` which is not parsed
  -- correctly in older versions of LLVM assembler (llvm-project#80571)
  -- So we use an `x` s.t. w64ToBase62 x > "R" > "L" > "r" > "l"
  , supply = 54
  }

renameDetUniq :: Unique -> DetRnM Unique
renameDetUniq uq = do
  m <- gets mapping
  case lookupUFM m uq of
    Nothing -> do
      new_w <- gets supply -- New deterministic unique in this `DetRnM`
      let --(_tag, _) = unpkUnique uq
          det_uniq = mkUnique 'Q' new_w
      modify (\DetUniqFM{mapping, supply} ->
        -- Update supply and mapping
        DetUniqFM
          { mapping = addToUFM mapping uq det_uniq
          , supply = supply + 1
          })
      return det_uniq
    Just det_uniq ->
      return det_uniq

-- The most important function here, which does the actual renaming.
detRenameCLabel :: CLabel -> DetRnM CLabel
detRenameCLabel = mapInternalNonDetUniques renameDetUniq

-- | We want to rename uniques in Ids, but ONLY internal ones.
detRenameId :: Id -> DetRnM Id
detRenameId i
  | isExternalName (varName i) = return i
  | otherwise = setIdUnique i <$> renameDetUniq (getUnique i)

-- | Similar to `detRenameId`, but for `Name`.
detRenameName :: Name -> DetRnM Name
detRenameName n
  | isExternalName n = return n
  | otherwise = setNameUnique n <$> renameDetUniq (getUnique n)

detRenameCmmGroup :: DetUniqFM -> DCmmGroup -> (DetUniqFM, CmmGroup)
detRenameCmmGroup dufm group = swap (runState (mapM detRenameCmmDecl group) dufm)
  where
    detRenameCmmDecl :: DCmmDecl -> DetRnM CmmDecl
    detRenameCmmDecl (CmmProc h lbl regs g)
      = do
        h' <- detRenameCmmTop h
        lbl' <- detRenameCLabel lbl
        regs' <- mapM detRenameGlobalRegUse regs
        g' <- detRenameCmmGraph g
        return (CmmProc h' lbl' regs' g')
    detRenameCmmDecl (CmmData sec d)
      = CmmData <$> detRenameSection sec <*> detRenameCmmStatics d

    detRenameCmmTop :: DCmmTopInfo -> DetRnM CmmTopInfo
    detRenameCmmTop (TopInfo (DWrap i) b)
      = TopInfo . mapFromList <$> mapM (detRenamePair detRenameLabel detRenameCmmInfoTable) i <*> pure b

    detRenameCmmGraph :: DCmmGraph -> DetRnM CmmGraph
    detRenameCmmGraph (CmmGraph entry bs)
      = CmmGraph <$> detRenameLabel entry <*> detRenameGraph bs

    detRenameGraph = \case
      GNil  -> pure GNil
      GUnit block -> GUnit <$> detRenameBlock block
      GMany m1 b m2 -> GMany <$> detRenameMaybeBlock m1 <*> detRenameBody b <*> detRenameMaybeBlock m2

    detRenameBody (DWrap b)
      = mapFromList <$> mapM (detRenamePair detRenameLabel detRenameBlock) b

    detRenameCmmStatics :: CmmStatics -> DetRnM CmmStatics
    detRenameCmmStatics
      (CmmStatics clbl info ccs lits1 lits2)
        = CmmStatics <$> detRenameCLabel clbl <*> detRenameCmmInfoTable info <*> pure ccs <*> mapM detRenameCmmLit lits1 <*> mapM detRenameCmmLit lits2
    detRenameCmmStatics
      (CmmStaticsRaw lbl sts)
        = CmmStaticsRaw <$> detRenameCLabel lbl <*> mapM detRenameCmmStatic sts

    detRenameCmmInfoTable :: CmmInfoTable -> DetRnM CmmInfoTable
    detRenameCmmInfoTable
      CmmInfoTable{cit_lbl, cit_rep, cit_prof, cit_srt, cit_clo}
        = CmmInfoTable <$> detRenameCLabel cit_lbl <*> pure cit_rep <*> pure cit_prof <*> detRenameMaybe detRenameCLabel cit_srt <*>
           (case cit_clo of
              Nothing -> pure Nothing
              Just (an_id, ccs) -> Just . (,ccs) <$> detRenameId an_id)

    detRenameCmmStatic :: CmmStatic -> DetRnM CmmStatic
    detRenameCmmStatic = \case
      CmmStaticLit l -> CmmStaticLit <$> detRenameCmmLit l
      CmmUninitialised x -> pure $ CmmUninitialised x
      CmmString x -> pure $ CmmString x
      CmmFileEmbed f i -> pure $ CmmFileEmbed f i

    detRenameCmmLit :: CmmLit -> DetRnM CmmLit
    detRenameCmmLit = \case
      CmmInt i w -> pure $ CmmInt i w
      CmmFloat r w -> pure $ CmmFloat r w
      CmmVec lits -> CmmVec <$> mapM detRenameCmmLit lits
      CmmLabel lbl -> CmmLabel <$> detRenameCLabel lbl
      CmmLabelOff lbl i -> CmmLabelOff <$> detRenameCLabel lbl <*> pure i
      CmmLabelDiffOff lbl1 lbl2 i w ->
        CmmLabelDiffOff <$> detRenameCLabel lbl1 <*> detRenameCLabel lbl2 <*> pure i <*> pure w
      CmmBlock bid -> CmmBlock <$> detRenameLabel bid
      CmmHighStackMark -> pure CmmHighStackMark

    detRenameMaybeBlock :: MaybeO n (Block CmmNode a b) -> DetRnM (MaybeO n (Block CmmNode a b))
    detRenameMaybeBlock (JustO x) = JustO <$> detRenameBlock x
    detRenameMaybeBlock NothingO = pure NothingO

    detRenameBlock :: Block CmmNode n m -> DetRnM (Block CmmNode n m)
    detRenameBlock = \case
      BlockCO n bn -> BlockCO <$> detRenameCmmNode n <*> detRenameBlock bn
      BlockCC n1 bn n2 -> BlockCC <$> detRenameCmmNode n1 <*> detRenameBlock bn <*> detRenameCmmNode n2
      BlockOC bn n -> BlockOC <$> detRenameBlock bn <*> detRenameCmmNode n
      BNil    -> pure BNil
      BMiddle n -> BMiddle <$> detRenameCmmNode n
      BCat    b1 b2 -> BCat <$> detRenameBlock b1 <*> detRenameBlock b2
      BSnoc   bn n -> BSnoc <$> detRenameBlock bn <*> detRenameCmmNode n
      BCons   n bn -> BCons <$> detRenameCmmNode n <*> detRenameBlock bn

    detRenameCmmNode :: CmmNode n m -> DetRnM (CmmNode n m)
    detRenameCmmNode = \case
      CmmEntry l t -> CmmEntry <$> detRenameLabel l <*> detRenameCmmTick t
      CmmComment fs -> pure $ CmmComment fs
      CmmTick tickish -> pure $ CmmTick tickish
      CmmUnwind xs -> CmmUnwind <$> mapM (detRenamePair detRenameGlobalReg (detRenameMaybe detRenameCmmExpr)) xs
      CmmAssign reg e -> CmmAssign <$> detRenameCmmReg reg <*> detRenameCmmExpr e
      CmmStore e1 e2 align -> CmmStore <$> detRenameCmmExpr e1 <*> detRenameCmmExpr e2 <*> pure align
      CmmUnsafeForeignCall ftgt cmmformal cmmactual ->
        CmmUnsafeForeignCall <$> detRenameForeignTarget ftgt <*> mapM detRenameLocalReg cmmformal <*> mapM detRenameCmmExpr cmmactual
      CmmBranch l -> CmmBranch <$> detRenameLabel l
      CmmCondBranch pred t f likely ->
        CmmCondBranch <$> detRenameCmmExpr pred <*> detRenameLabel t <*> detRenameLabel f <*> pure likely
      CmmSwitch e sts -> CmmSwitch <$> detRenameCmmExpr e <*> mapSwitchTargetsA detRenameLabel sts
      CmmCall tgt cont regs args retargs retoff ->
        CmmCall <$> detRenameCmmExpr tgt <*> detRenameMaybe detRenameLabel cont <*> mapM detRenameGlobalRegUse regs
                <*> pure args <*> pure retargs <*> pure retoff
      CmmForeignCall tgt res args succ retargs retoff intrbl ->
        CmmForeignCall <$> detRenameForeignTarget tgt <*> mapM detRenameLocalReg res <*> mapM detRenameCmmExpr args
                       <*> detRenameLabel succ <*> pure retargs <*> pure retoff <*> pure intrbl

    detRenameCmmExpr :: CmmExpr -> DetRnM CmmExpr
    detRenameCmmExpr = \case
      CmmLit l -> CmmLit <$> detRenameCmmLit l
      CmmLoad e t a -> CmmLoad <$> detRenameCmmExpr e <*> pure t <*> pure a
      CmmReg r -> CmmReg <$> detRenameCmmReg r
      CmmMachOp mop es -> CmmMachOp mop <$> mapM detRenameCmmExpr es
      CmmStackSlot a i -> CmmStackSlot <$> detRenameArea a <*> pure i
      CmmRegOff r i -> CmmRegOff <$> detRenameCmmReg r <*> pure i

    detRenameForeignTarget :: ForeignTarget -> DetRnM ForeignTarget
    detRenameForeignTarget = \case
        ForeignTarget e fc -> ForeignTarget <$> detRenameCmmExpr e <*> pure fc
        PrimTarget cmop -> pure $ PrimTarget cmop

    detRenameArea :: Area -> DetRnM Area
    detRenameArea Old = pure Old
    detRenameArea (Young l) = Young <$> detRenameLabel l

    detRenameLabel :: Label -> DetRnM Label
    detRenameLabel lbl
      = mkHooplLabel . getKey <$> renameDetUniq (getUnique lbl)

    detRenameSection :: Section -> DetRnM Section
    detRenameSection (Section ty lbl)
      = Section ty <$> detRenameCLabel lbl

    detRenameCmmReg :: CmmReg -> DetRnM CmmReg
    detRenameCmmReg = \case
      CmmLocal l -> CmmLocal <$> detRenameLocalReg l
      CmmGlobal x -> pure $ CmmGlobal x

    detRenameLocalReg :: LocalReg -> DetRnM LocalReg
    detRenameLocalReg (LocalReg uq t)
      = LocalReg <$> renameDetUniq uq <*> pure t

    -- Global registers don't need to be renamed.
    detRenameGlobalReg :: GlobalReg -> DetRnM GlobalReg
    detRenameGlobalReg = pure
    detRenameGlobalRegUse :: GlobalRegUse -> DetRnM GlobalRegUse
    detRenameGlobalRegUse = pure

    -- todo: We may have to change this to get deterministic objects with ticks.
    detRenameCmmTick :: CmmTickScope -> DetRnM CmmTickScope
    detRenameCmmTick = pure

    detRenameMaybe _ Nothing = pure Nothing
    detRenameMaybe f (Just x) = Just <$> f x

    detRenamePair f g (a, b) = (,) <$> f a <*> g b

detRenameIPEMap :: DetUniqFM -> InfoTableProvMap -> (DetUniqFM, InfoTableProvMap)
detRenameIPEMap dufm InfoTableProvMap{ provDC, provClosure, provInfoTables } =
    (dufm2, InfoTableProvMap { provDC, provClosure = cm, provInfoTables })
  where
    (cm, dufm2) = runState (detRenameClosureMap provClosure) dufm

    detRenameClosureMap :: ClosureMap -> DetRnM ClosureMap
    detRenameClosureMap m =
      -- `eltsUDFM` preserves the deterministic order, but it doesn't matter
      -- since we will rename all uniques deterministically, thus the
      -- reconstructed map will necessarily be deterministic too.
      listToUDFM <$> mapM (\(n,r) -> do
        n' <- detRenameName n
        return (n', (n', r))
        ) (eltsUDFM m)