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.
  -}
  
  {-# LANGUAGE TypeApplications #-}
  module Derive.HackDeclarationTarget
    ( deriveHackDeclarationTarget
    ) where
  
  import Prelude hiding (span)
  
  import Control.Concurrent (getNumCapabilities)
  import Control.Concurrent.Async
  import Control.Concurrent.STM
  import Control.DeepSeq
  import Control.Exception
  import Control.Monad
  import qualified Data.Int as Int
  import qualified Data.IntervalMap.Generic.Strict as IM
  import Data.Maybe
  import Data.List (foldl')
  import Data.List.Extra (chunksOf)
  import Util.Log (logInfo)
  
  import Glean
  import Glean.Angle as Angle
  import qualified Glean.Schema.Hack.Types as Hack
  import qualified Glean.Schema.Src.Types as Src
  import Glean.Util.PredMap (PredMap)
  import qualified Glean.Util.PredMap as PredMap
  import Glean.Util.Mutex
  
  import Derive.Env
  import Derive.Types
  
  type TargetSpans = IM.IntervalMap ByteSpan [Hack.Declaration]
  type Files = PredMap Src.File ()
  type Sources = PredMap Src.File [(Hack.Declaration, ByteSpan)]
  type Targets = PredMap Src.File [[TargetSpans]]
  
  newtype ByteSpan = ByteSpan Src.ByteSpan
    deriving (Eq, Ord, Show)
  
  instance IM.Interval ByteSpan Int.Int64 where
    lowerBound (ByteSpan bs) = unNat $ Src.byteSpan_start bs
    upperBound (ByteSpan bs) =
      unNat (Src.byteSpan_start bs) +
      unNat (Src.byteSpan_length bs)
    leftClosed _ = True -- The interval contains the start value
    rightClosed _ = False -- The interval does not contain start + length
  
  getFilesQuery :: Query Hack.FileXRefs
  getFilesQuery = Angle.query $ Angle.predicate @Hack.FileXRefs Angle.wild
  
  getSourcesQuery :: [IdOf Src.File] -> Query Hack.DeclarationSpan
  getSourcesQuery files =
    Angle.query $ Angle.vars $ \x file decls decl -> x `Angle.where_`
    [ file .= Angle.elementsOf (Angle.array $ map Angle.factId files)
    , wild .= Angle.predicate @Hack.FileDeclarations  (
        rec $
          field @"file" (asPredicate file) $
          field @"declarations" decls
        end
      )
    , decl .= Angle.elementsOf decls
    , x .= Angle.predicate @Hack.DeclarationSpan (
        rec $
          field @"declaration" decl $
          field @"file" (asPredicate file)
        end
      )
    ]
  
  getTargetsQuery :: [IdOf Src.File] ->  Query Hack.FileXRefs
  getTargetsQuery files =
    Angle.query $ Angle.vars $ \x file -> x `Angle.where_`
    [ file .= Angle.elementsOf (Angle.array $ map Angle.factId files)
    , x .= Angle.predicate @Hack.FileXRefs (
        rec $
          field @"file" (asPredicate file)
        end
      )
    ]
  
  getFiles
    :: (Backend e)
    => e
    -> Config
    -> IO Files
  getFiles e cfg =
    runQueryEach e (cfgRepo cfg) keysQuery mempty $ \ pm p -> do
      return $! PredMap.insert (getId $ Hack.fileXRefs_key_file p) () pm
    where
        keysQuery = limitBytes (cfgMaxQuerySize cfg) $ keys getFilesQuery
  
  loadPredMap
    :: (Backend e, Predicate p)
    => e
    -> Config
    -> Query p
    -> (KeyType p -> (IdOf k, v))
    -> IO (PredMap k [v])
  loadPredMap e cfg q f =
    runQueryEach e (cfgRepo cfg) keysQuery mempty $ \ pm p -> do
      let (k, v) = f p
      return $! PredMap.alter (Just . maybe [v] (v :)) k pm
      where
        keysQuery = limitBytes (cfgMaxQuerySize cfg) $ keys q
  
  getSources :: Backend e => e -> Config -> [IdOf Src.File] -> IO Sources
  getSources e cfg files = do
    loadPredMap e cfg (getSourcesQuery files) f
    where
      f key = (getId file, (declaration, ByteSpan span))
        where
          file = Hack.declarationSpan_key_file key
          declaration = Hack.declarationSpan_key_declaration key
          span = Hack.declarationSpan_key_span key
  
  getTargets :: Backend e => e -> Config -> [IdOf Src.File] -> IO Targets
  getTargets e cfg files = do
    loadPredMap e cfg (getTargetsQuery files) f
    where
      f key = (getId file, targets)
        where
          file = Hack.fileXRefs_key_file key
          xrefs = Hack.fileXRefs_key_xrefs key
          targets = mapMaybe target xrefs
          target xref = case Hack.xRef_target xref of
            Hack.XRefTarget_declaration decl ->
              let spans = Hack.xRef_ranges xref
                  (_, im) = foldl' (buildIntervalMap decl) (0, IM.empty) spans
              in Just im
            Hack.XRefTarget_occurrence _ -> Nothing
            Hack.XRefTarget_EMPTY -> Nothing
          buildIntervalMap
            :: Hack.Declaration
            -> (Int.Int64, TargetSpans)
            -> Src.RelByteSpan
            -> (Int.Int64, TargetSpans)
          buildIntervalMap decl (offset, im) span =
              (nextOffset, IM.insert byteSpan [decl] im)
            where
              (nextOffset, byteSpan) = relToAbsolute offset span
          relToAbsolute offset rel =
            (start, ByteSpan $ Src.ByteSpan (Nat start) len)
            where
              start = unNat (Src.relByteSpan_offset rel) + offset
              len = Src.relByteSpan_length rel
  
  declarationTargetsForFile
    :: (Sources, Targets)
    -> IdOf Src.File
    -> [Hack.DeclarationTarget_key]
  declarationTargetsForFile (sources, targets) fileId =
    declTargets
    where
      sourceList = PredMap.findWithDefault [] fileId sources
      targetList = PredMap.findWithDefault [] fileId targets
      targetIm = IM.unionsWith (++) $ concat targetList
      declTargets = concatMap (declarationTargetsForSource targetIm) sourceList
  
  declarationTargetsForSource
    :: TargetSpans
    -> (Hack.Declaration, ByteSpan)
    -> [Hack.DeclarationTarget_key]
  declarationTargetsForSource targetSpans (source, sourceSpan) =
    map (Hack.DeclarationTarget_key source) targets
    where
      targets = concatMap snd $ IM.toList targetsIm
      targetsIm = targetSpans `IM.within` sourceSpan
  
  writeHackDeclarationTargets
    :: Writer -> TBQueue (Maybe [Hack.DeclarationTarget_key]) -> IO ()
  writeHackDeclarationTargets writer tbqueue = do
    next <- atomically $ readTBQueue tbqueue
    case next of
      Nothing -> return ()
      Just declTargets -> do
        forM_ declTargets $ \declTarget -> writeFacts writer $
          makeFact_ @Hack.DeclarationTarget declTarget
        writeHackDeclarationTargets writer tbqueue
  
  
  writeChunkedParallelMap
    :: (NFData b)
    => Int
    -> TBQueue (Maybe [b])
    -> ((Int,[a]) -> IO ctx)
    -> (ctx -> a -> [b])
    -> [a]
    -> IO ()
  writeChunkedParallelMap nchunks tbqueue chunkCtx f xs = do
    logInfo $ "calculating " ++ show nchunks ++ " chunks for "
      ++ show (length xs) ++ " inputs"
    mutex <- newMutex ()
    forConcurrently_ (zip [1..] chunks) $ \(i, chunk) -> do
      -- Use mutex to avoid many queries in flight at once
      ctx <- withMutex_ mutex $ chunkCtx (i, chunk)
      ys <- evaluate $ force $ Just $ concatMap (f ctx) chunk
      atomically $ writeTBQueue tbqueue ys
      logInfo $ "calculated chunk " ++ show (i::Int)
    where
      chunks = chunksOf (max 1 chunkSize) xs
      chunkSize = (len + nchunks - 1) `div` nchunks
      len = length xs
  
  deriveHackDeclarationTarget
    :: Env
    -> IO ()
  deriveHackDeclarationTarget env@Env{..} = do
    numCapabilities <- getNumCapabilities
    logInfo "deriving HackDeclarationTarget"
    files <- getFiles envBackend envConfig
    logInfo "Got files"
    tbqueue :: TBQueue (Maybe [Hack.DeclarationTarget_key]) <-
      atomically $ newTBQueue (fromIntegral $ cfgMaxQueueSize envConfig)
    if null files
    then do
      logInfo "no file facts found - skipping"
      return ()
    else do
      let
        contextFor (i, fileIds) = do
          (sources, targets) <- concurrently
             (getSources envBackend envConfig fileIds
                >>= done ("Chunk " <> show i <> " got sources"))
             (getTargets envBackend envConfig fileIds
                >>= done ("Chunk " <> show i <> " got targets"))
          return (sources, targets)
  
        calcDeclTargets = do
          writeChunkedParallelMap
            nchunks
            tbqueue
            contextFor
            declarationTargetsForFile
            (PredMap.keys files)
          atomically $ writeTBQueue tbqueue Nothing
          where
            nchunks = max 1 $ numCapabilities * cfgChunksPerCapability envConfig
  
        writeDeclTargets =
          withEnvWriter env $ \writer ->
            writeHackDeclarationTargets writer tbqueue
  
      ((), ()) <- runConcurrently $ (,)
        <$> Concurrently
             ( calcDeclTargets >>= done "calculated declaration targets")
        <*> Concurrently
              (writeDeclTargets >>= done "written facts")
      return ()
    where
      done msg x = logInfo msg >> return x
  