content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Glean.RTS.Bytecode.Code
    ( Code
    , Register
    , Label
    , CodeGen
    , OutputSupply
    , Many(..)
    , Optimised(..)
    , Meta(..)
    , literal
    , label
    , issue
    , issueEndBlock
    , constant
    , local
    , output
    , outputUninitialized
    , generate
    , castRegister
    , callSite
    , calledFrom
    , fullScan
    , vlog
    ) where
  
  import Control.Exception (assert)
  import Control.Monad
  import Control.Monad.Fix (MonadFix(..))
  import Control.Monad.ST (ST, runST)
  import Control.Monad.Trans
  import qualified Control.Monad.Trans.State.Strict as S
  import Data.Bits
  import Data.ByteString (ByteString)
  import Data.Functor
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import Data.IntMap (IntMap)
  import qualified Data.IntMap as IntMap
  import Data.IntSet (IntSet)
  import qualified Data.IntSet as IntSet
  import Data.List (mapAccumL, sortBy)
  import Data.Maybe
  import Data.Ord
  import qualified Data.Vector as V
  import qualified Data.Vector.Primitive as VP
  import qualified Data.Vector.Primitive.Mutable as VPM
  import qualified Data.Vector.Storable as VS
  import Data.Word (Word64)
  import GHC.Stack
  import qualified Util.Log as Log
  
  import Glean.Bytecode.Types
  import Glean.RTS.Types (Pid)
  import Glean.RTS.Bytecode.Gen.Instruction
  import Glean.RTS.Bytecode.Supply
  import Glean.RTS.Foreign.Bytecode (Subroutine, subroutine)
  
  -- | A basic block
  newtype Block = Block
    { -- | Instructions (reversed)
      blockInsns :: [Insn]
    }
  
  data CodeS = CodeS
    { -- | Label of current block
      csLabel :: {-# UNPACK #-} !Label
  
      -- | Instructions in current block (reversed)
    , csInsns :: [Insn]
  
      -- | Blocks produced so far (reversed)
    , csBlocks :: [Block]
  
      -- | All constant values used by the subroutine. These will be preloaded
      -- into registers at the start.
    , csConstants :: [Word64]
  
      -- | Cached `length csConstants`
    , csNextConstant :: !(Register 'Word)
  
      -- | Map known constants to their registers
    , csConstantMap :: IntMap (Register 'Word)
  
      -- | All literals in the subroutine.
    , csLiterals :: HashMap ByteString Word64
  
      -- | Cached `length csLiterals`
    , csLiteralsSize :: !Word64
  
      -- | Currently used number of local registers
    , csNextLocal :: {-# UNPACK #-} !(Register 'Word)
  
      -- | Maximum number of local registers used so far
    , csMaxLocal :: {-# UNPACK #-} !(Register 'Word)
  
      -- | Currently used number of binary::Output registers
    , csNextOutput :: {-# UNPACK #-} !(Register 'BinaryOutputPtr)
  
      -- | Maximum number of binary::Output registers
    , csMaxOutputs :: {-# UNPACK #-} !(Register 'BinaryOutputPtr)
  
      -- | Predicates we perform full scans on.
      -- Repeated entries mean multiple scans of the same predicate
    , csFullScans:: [Pid]
    }
  
  -- | Code gen monad
  newtype Code a = Code { runCode :: S.StateT CodeS IO a }
    deriving(Functor, Applicative, Monad, MonadFix)
  
  -- | Things that generate code of the form
  -- > Register t1 -> ... -> Register tn -> Code a
  --
  class CodeGen s cg where
    genCode :: cg -> S.State s (Code (CodeResult cg))
  
  instance CodeGen s (Code a) where
    genCode = pure
  
  instance (Supply a s, CodeGen s cg) => CodeGen s (a -> cg) where
    genCode f = S.state supply >>= genCode . f
  
  -- | Allocate `n` registers where `n` isn't statically known.
  -- Example:
  --
  -- local $ Many n $ \regs -> ...
  data Many a cg = Many Int ([a] -> cg)
  
  instance (Supply a s, CodeGen s cg) => CodeGen s (Many a cg) where
    genCode (Many n f) = replicateM n (S.state supply) >>= genCode . f
  
  type family CodeResult cg
  type instance CodeResult (Code a) = a
  type instance CodeResult (a -> cg) = CodeResult cg
  type instance CodeResult (Many a cg) = CodeResult cg
  
  -- | Load a constant value into a register. This will happen at the start of
  -- the subroutine, effectively giving us a poor man's version of constant
  -- hoisting.
  constant :: Word64 -> Code (Register 'Word)
  constant w = Code $ do
    s@CodeS{..} <- S.get
    case IntMap.lookup (fromIntegral w) csConstantMap of
      Just r -> return r
      Nothing -> do
        S.put s
          { csConstants = w : csConstants
          , csNextConstant = succ csNextConstant
          , csConstantMap =
              IntMap.insert (fromIntegral w) csNextConstant csConstantMap
          }
        return csNextConstant
  
  -- | Generate a chunk of code with reserved fresh registers. The registers can
  -- be reused afterwards.
  local :: CodeGen RegSupply cg => cg -> Code (CodeResult cg)
  local = runLocal
    regSupply
    csNextLocal
    (\r s -> s { csNextLocal = r })
    (\r s -> s { csMaxLocal = max (csMaxLocal s) r })
  
  -- | Supply of @Register 'BinaryOutputPtr@ (cf. 'output').
  newtype OutputSupply = OutputSupply RegSupply
    deriving (Supply (Register 'BinaryOutputPtr))
  
  -- | Declare registers of type @Register 'BinaryOutputPtr@.  The registers
  -- are inputs to the subroutine. Use 'resetOutput' to reset the byte array
  -- stored in these registers to empty.
  outputUninitialized :: CodeGen OutputSupply cg => cg -> Code (CodeResult cg)
  outputUninitialized = runLocal
    (OutputSupply . regSupply)
    csNextOutput
    (\r s -> s { csNextOutput = r })
    (\r s -> s { csMaxOutputs = max (csMaxOutputs s) r })
  
  output :: (Register 'BinaryOutputPtr -> Code a) -> Code a
  output f = outputUninitialized $ \out -> do
    issue $ ResetOutput out
    f out
  
  runLocal
    :: (Supply r s, CodeGen s cg)
    => (r -> s)
    -> (CodeS -> r)
    -> (r -> CodeS -> CodeS)
    -> (r -> CodeS -> CodeS)
    -> cg
    -> Code (CodeResult cg)
  runLocal make get set setmax cg = do
    r <- Code $ S.gets get
    let (gen, sup) = S.runState (genCode cg) $ make r
        !next = peekSupply sup
    Code $ S.modify' $ setmax next . set next
    x <- gen
    Code $ S.modify' $ set r
    return x
  
  -- | Poor man's function calls
  --
  -- * Put the label of the return address into a register with
  --   'loadReg', jump to the code, and return with 'jumpReg'.
  --
  -- * We have to avoid local registers at the call site(s) clashing
  --   with local registers in the called code.  So 'callSite' remembers
  --   the number of locals in scope at the call site(s) and
  --   'calledCode' uses the high-water mark of the call sites as the
  --   base for its local registers.
  --
  data CallSite = CallSite
    { callSiteNextLocal :: Register 'Word
    , callSiteNextOutput :: Register 'BinaryOutputPtr
    }
  
  callSite :: Code CallSite
  callSite = do
    CodeS{..} <- Code S.get
    return (CallSite csNextLocal csNextOutput)
  
  calledFrom :: [CallSite] -> Code a -> Code a
  calledFrom frames inner = do
    CodeS{..} <- Code S.get
    Code $ S.modify' $ \s -> s
      { csNextLocal = maximum (map callSiteNextLocal frames)
      , csNextOutput = maximum (map callSiteNextOutput frames) }
    x <- inner
    Code $ S.modify' $ \s -> s
      { csNextLocal = csNextLocal
      , csNextOutput = csNextOutput }
    return x
  
  -- | Register that a query statement performs a full scan over a predicate.
  fullScan :: Pid -> Code ()
  fullScan pid = Code $ S.modify' $ \s -> s { csFullScans = pid : csFullScans s }
  
  data Optimised = Optimised | Unoptimised
    deriving(Eq,Ord,Enum,Bounded,Show)
  
  -- | Metadata about the subroutine
  newtype Meta = Meta
    { meta_fullScans :: [Pid]
    }
  
  -- | Generate a 'Subroutine', allocating input registers as necessary.
  -- Example:
  --
  -- generate $ \reg1 reg2 -> do
  --   add reg1 reg2 reg1
  --   ret
  --
  generate
    :: (CodeGen RegSupply cg, CodeResult cg ~ ())
    => Optimised -> cg -> IO (Meta, Subroutine t)
  generate opt cg = do
    let (gen, sup) = S.runState (genCode cg) $ regSupply $ register Input 0
        !nextInput = peekSupply sup
    ((), CodeS{..}) <- S.runStateT (runCode gen) CodeS
          { csLabel = Label 0
          , csInsns = []
          , csBlocks = []
          , csConstants = []
          , csConstantMap = IntMap.empty
          , csNextConstant = register Constant 0
          , csLiterals = HashMap.empty
          , csLiteralsSize = 0
          , csNextLocal = register Local 0
          , csMaxLocal = register Local 0
          , csNextOutput = castRegister nextInput
          , csMaxOutputs = castRegister nextInput
          , csFullScans = mempty }
        -- sanity check
    when (not $ null csInsns) $ fail "unterminated basic block"
    let -- output registers go after input registers
        finalInputSize = registerIndex csMaxOutputs
        constantsSize = registerIndex csNextConstant
        get_label pc label =
          let addr = offsets VP.! fromLabel label
          in assert (addr /= maxBound) $ addr - pc
        get_reg :: forall ty . Register ty -> Word64
        get_reg r = case registerSegment r of
          Input -> assert (n < finalInputSize) n
          Constant -> assert (n < constantsSize) (n + finalInputSize)
          Local -> n + finalInputSize + constantsSize
          where
            !n = registerIndex r
        optimise = case opt of
          Optimised -> shortcut
          Unoptimised -> id
        (insns, offsets) = layout $ optimise CFG
          { cfgBlocks = V.fromListN (fromLabel csLabel) $ reverse csBlocks
          , cfgEntry = Label 0
          }
        code = concat $ snd $ mapAccumL
          (\offset insn ->
            let !next = offset + insnSize insn
            in
            (next, insnWords get_reg (get_label next) insn))
          0
          insns
        meta = Meta csFullScans
    (meta,) <$> subroutine
      (VS.fromListN (length code) code)
      finalInputSize
      (finalInputSize - registerIndex nextInput)
      (registerIndex csMaxLocal + constantsSize)
      (reverse csConstants)
      (map fst $ sortBy (comparing snd) $ HashMap.toList csLiterals)
  
  
  -- | Control flow graph
  data CFG = CFG
    { -- | Basic blocks
      cfgBlocks :: !(V.Vector Block)
  
      -- | Entry block
    , cfgEntry :: {-# UNPACK #-} !Label
    }
  
  -- | Inline blocks containing only one instruction and short-circuit labels
  -- which point to blocks consisting of a single Jump.
  --
  -- NOTE: This will leave behind unreachable blocks.
  shortcut :: CFG -> CFG
  shortcut cfg@CFG{..}
    | V.all isNothing shortcuts = cfg
    | otherwise = CFG
        { cfgBlocks = V.imap
            (\i Block{..} ->
                Block { blockInsns = mapLabels relabel <$> inline i blockInsns })
            cfgBlocks
        , cfgEntry = relabel cfgEntry
        }
    where
      inline i _
        | Just insn <- shortcuts V.! i = [insn]
      inline _ (Jump target : insns)
        | Just insn <- shortcuts V.! fromLabel target = insn : insns
      inline _ insns = insns
  
      -- FIXME: This can loop if the generated code contains infinite loops
      shortcuts = cfgBlocks <&> \Block{..} -> case blockInsns of
        [insn] -> Just $ case insn of
          Jump target | Just insn' <- shortcuts V.! fromLabel target -> insn'
          _ -> insn
        _ -> Nothing
  
      relabel label
        | Just (Jump target) <- shortcuts V.! fromLabel label = target
        | otherwise = label
  
  data Layout s = Layout
    { -- | Current offset in the instruction stream
      layoutOffset :: {-# UNPACK #-} !Word64
  
      -- | Label offsets, 'maxBound' for blocks which haven't been emitted yet
    , layoutLabels :: !(VPM.STVector s Word64)
  
      -- | Insn stream (all reversed)
    , layoutInsns :: [[Insn]]
  
      -- | Blocks we want to emit (some of them might have been emitted already)
    , layoutTodo :: !IntSet
    }
  
  -- | Compute a flat instruction stream for a subroutine as well as a mapping
  -- from labels to their offsets in that stream. Note that we don't emit
  -- unreachable blocks - a block that hasn't been emitted because it's dead will
  -- have the magic value of maxBound in the label->offset mapping.
  --
  -- This is really simple at the moment.
  --
  -- * Start with the entry block.
  -- * If the current block ends with an unconditional jump and the block it jumps
  --   to hasn't been emitted yet, continue with that block, thus saving the jump.
  -- * Otherwise, continue with the unemitted block with lowest label number.
  --
  -- There is obviously ample room for improvement here.
  layout :: CFG -> ([Insn], VP.Vector Word64)
  layout CFG{..} = runST $ do
    mlabels <- VPM.replicate (V.length cfgBlocks) maxBound
    emit cfgEntry Layout
      { layoutOffset = 0
      , layoutLabels = mlabels
      , layoutInsns = []
      , layoutTodo = IntSet.empty
      }
    where
      -- Emit a specific block which must not have been emitted already
      emit :: Label -> Layout s -> ST s ([Insn], VP.Vector Word64)
      emit !label layout@Layout{..} = do
        VPM.write layoutLabels (fromLabel label) layoutOffset
  
        case blockInsns (cfgBlocks V.! fromLabel label) of
          -- Handle unconditional jumps specially
          Jump target : insns -> do
            s <- stillTodo layout target
            if s
              then emit target $ addInsns insns layout
              else emitNext $ addInsns (Jump target : insns) layout
  
          insns -> emitNext $ addInsns insns layout
  
      -- Emit the numerically lowest unemitted block (if any)
      emitNext :: Layout s -> ST s ([Insn], VP.Vector Word64)
      emitNext layout@Layout{..}
        | Just (label, todo) <- IntSet.minView layoutTodo = do
            s <- stillTodo layout $ Label label
            (if s then emit (Label label) else emitNext) layout{layoutTodo = todo}
        | otherwise = do
            labels <- VP.unsafeFreeze layoutLabels
            return (reverse $ concat layoutInsns, labels)
  
      -- Has a block already been emitted
      stillTodo :: Layout s -> Label -> ST s Bool
      stillTodo Layout{..} label =
        (== maxBound) <$> VPM.read layoutLabels (fromLabel label)
  
      -- Emit instructions
      addInsns :: [Insn] -> Layout s -> Layout s
      addInsns insns !layout = layout
        { layoutOffset = layoutOffset layout + sum (map insnSize insns)
        , layoutTodo = foldr
            (IntSet.insert . fromLabel)
            (layoutTodo layout)
            (concatMap insnLabels insns)
        , layoutInsns = insns : layoutInsns layout
        }
  
  -- | Frame segment which a register belongs to
  data Segment = Input | Constant | Local
    deriving(Eq,Ord,Enum,Bounded,Show)
  
  -- A register has a 'Segment' and an index within the segment, packed into
  -- a word.
  
  -- | Create a register with the given 'Segment' and index
  register :: Segment -> Word64 -> Register a
  register s i = Register $ (fromIntegral (fromEnum s) `shiftL` 62) .|. i
  
  -- | Get the 'Segment' of the register
  registerSegment :: Register a  -> Segment
  registerSegment (Register i) = toEnum $ fromIntegral (i `shiftR` 62)
  
  -- | Get the index of the register within its segment
  registerIndex :: Register a -> Word64
  registerIndex (Register i) = i .&. 0x3FFFFFFFFFFFFFFF
  
  -- | Start a new basic block. The previous block will be terminated by the
  -- supplied unconditional jump instruction or, if none is provided, by a jump
  -- to the new block.
  newBlock :: Maybe Insn -> Code ()
  newBlock terminator = Code $ do
    s@CodeS{..} <- S.get
    let !label = succ csLabel
        insn = fromMaybe (Jump label) terminator
    S.put $! s
      { csLabel = label
      , csInsns = []
      , csBlocks = Block (insn : csInsns) : csBlocks
      }
  
  -- | Add a literal to the literal table and yield its index
  literal :: ByteString -> Code Literal
  literal lit = Code $ Literal . fromIntegral <$> do
    m <- S.gets csLiterals
    case HashMap.lookup lit m of
      Just n -> return n
      Nothing -> do
        n <- S.gets csLiteralsSize
        S.modify' $ \s@CodeS{..} -> s
          { csLiterals = HashMap.insert lit n csLiterals
          , csLiteralsSize = csLiteralsSize + 1 }
        return n
  
  -- | Yield a label for the current position in the code
  label :: Code Label
  label = do
    insns <- Code $ S.gets csInsns
    when (not $ null insns) $ newBlock Nothing
    Code $ S.gets csLabel
  
  -- | Issue an instruction which doesn't modify the program counter
  issue :: Insn -> Code ()
  issue insn = Code $ S.modify' $ \s@CodeS{..} -> s { csInsns = insn : csInsns }
  
  -- | Issue an instruction which always modifies the program counter
  issueEndBlock :: Insn -> Code ()
  issueEndBlock = newBlock . Just
  
  -- | Make some noise during code generation. Helpful for debugging.
  --
  vlog :: HasCallStack => String -> Code ()
  vlog msg = Code $ lift (Log.vlog 2 msg)
  