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.Common
    ( getFileXRefsFor
    , getIndirectTargets
    , resolve
    , Indirects
    , XRefs
    , emptyBatch
    , batchingYield
    ) where
  
  import Control.Monad
  import Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import Data.Maybe
  import qualified Data.Vector as V
  import qualified Data.Vector.Mutable as VM
  
  import Util.List
  
  import Glean
  import Glean.Angle
  import qualified Glean.Schema.Cxx1.Types as Cxx
  import Glean.Util.PredMap (PredMap)
  import qualified Glean.Util.PredMap as PredMap
  import Glean.Util.PredSet (PredSet)
  import qualified Glean.Util.PredSet as PredSet
  
  import Derive.Types
  
  -- | Keys are xRefIndirectTarget_id, values are xRefIndirectTarget_key_target
  -- Note: XRefTarget_indirect constructor does not appear in these values
  type Indirects = PredMap Cxx.XRefIndirectTarget Cxx.XRefTarget
  
  -- Keys are fileXRefMap_id, values are sets of targets for each external target
  -- in the map.
  type XRefs =
    PredMap
      Cxx.FileXRefMap
      (PredSet Cxx.FileXRefs, V.Vector (HashSet Cxx.XRefTarget))
  
  type MutableXRefs =
    PredMap
      Cxx.FileXRefMap
      (PredSet Cxx.FileXRefs, VM.IOVector (HashSet Cxx.XRefTarget))
  
  getIndirectTargets
    :: Backend e => e -> Config -> IO Indirects
  getIndirectTargets e cfg = do
  
    let q = maybe id limit (cfgMaxQueryFacts cfg) $
          limitBytes (cfgMaxQuerySize cfg) allFacts
  
    -- Build indirect map that is 1-level deep.  Chains of indirections not
    -- yet resolved so the values may be XRefTarget_indirect.
    targets <- runQueryEach e (cfgRepo cfg) q mempty $ \targets fact -> do
      return $ PredMap.insert
        (getId fact)
        (Cxx.xRefIndirectTarget_key_target $ fromJust $ getFactKey fact)
        targets
  
    let -- Use 'targets' to resolve chains of indirections to resolve
        -- (or remove) all XRefTarget_indirect constructors.
        squash (Cxx.XRefTarget_indirect x)
          | Just next <- PredMap.lookup (getId x) targets = squash next
          | otherwise = Nothing -- shouldn't happen
        squash x = Just x
  
    return $ PredMap.mapMaybe squash targets
  
  -- | The result never has XRefTarget_indirect constructor
  resolve :: Indirects -> Cxx.XRefTarget -> Maybe Cxx.XRefTarget
  resolve indirects (Cxx.XRefTarget_indirect x) =
    PredMap.lookup (getId x) indirects
  resolve _ x = Just x
  
  getFileXRefsFor
    :: Backend e
    => e
    -> [IdOf Cxx.FileXRefMap]
    -> Config
    -> IO XRefs
  getFileXRefsFor _ [] _ = return PredMap.empty
  getFileXRefsFor e xmapIds cfg = do
    let
      q :: Query Cxx.FileXRefs
      q = limitBytes (cfgMaxQuerySize cfg) $
            expanding @Cxx.XRefTargets $
            query $
            vars $ \(xmap :: Angle Cxx.FileXRefMap) xrefs ->
              xrefs `where_` [
                xmap .= elementsOf (factIdsArray xmapIds),
                xrefs .= (predicate @Cxx.FileXRefs $ rec $
                  field @"xmap" (asPredicate xmap) end)
              ]
  
    (_, xrefs) <- runQueryEachBatch e (cfgRepo cfg) q (mempty, mempty)
      withFileXRefs
    traverse (\(deps, xs) -> (deps,) <$> V.unsafeFreeze xs) xrefs
  
    where
    withFileXRefs
      :: (PredMap Cxx.XRefTargets (HashSet Cxx.XRefTarget), MutableXRefs)
      -> [Cxx.FileXRefs]
      -> IO (PredMap Cxx.XRefTargets (HashSet Cxx.XRefTarget), MutableXRefs)
    withFileXRefs (oldTargetMap, xrefs) fileXRefs = do
      let
        newTargets = uniq
            [ t
            | Cxx.FileXRefs _ (Just (Cxx.FileXRefs_key _ targets)) <- fileXRefs
            , t <- targets
            , not (getId t `PredMap.member` oldTargetMap)
            ]
  
        newTargetMap = PredMap.fromList
            [ (IdOf $ Fid i, HashSet.fromList k)
            | Cxx.XRefTargets i (Just k) <- newTargets
            ]
  
        targetMap = PredMap.union newTargetMap oldTargetMap
  
      let
        add xrefs (Cxx.FileXRefs i (Just (Cxx.FileXRefs_key xmap targetIds))) = do
          let !id = getId xmap
              targets = [ key | t <- targetIds,
                Just key <- [PredMap.lookup (getId t) targetMap]]
          case PredMap.lookup id xrefs of
            Just (deps, xs) -> do
              forM_ (zip [0 .. VM.length xs - 1] targets) $ \(i, ts) -> do
                x <- VM.unsafeRead xs i
                VM.unsafeWrite xs i $! HashSet.union ts x
              let !newDeps = PredSet.insert (IdOf $ Fid i) deps
              return $ PredMap.insert id (newDeps, xs) xrefs
            Nothing -> do
              xs <- VM.new (length targets)
              forM_ (zip [0..] targets) $ \(i, ts) ->
                VM.unsafeWrite xs i ts
              return $
                PredMap.insert id (PredSet.singleton (IdOf $ Fid i), xs) xrefs
        add xrefs _ = return xrefs
  
      newXrefs <- foldM add xrefs fileXRefs
      return (targetMap, newXrefs)
  
  -- -----------------------------------------------------------------------------
  -- Batching utils
  
  emptyBatch :: (Int, [a])
  emptyBatch = (0, [])
  
  batchingYield
    :: Int  -- batch size
    -> ([a] -> IO ())  -- yield a complete batch
    -> Int -- size of item
    -> a -- item
    -> (Int, [a]) -- current accumuulated batch
    -> IO (Int, [a]) -- new accumulated batch
  batchingYield batchSize yield n item (size, batch)
    | size+n >= batchSize = do yield (item:batch); return emptyBatch
    | otherwise = return (size+n, item:batch)
  