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.Query.UserQuery
    ( userQueryFacts
    , userQuery
    , userQueryWrites
    , schemaVersionForQuery
    ) where
  
  import Control.Applicative
  import Control.DeepSeq
  import Control.Exception
  import Control.Monad
  import Control.Monad.Except
  import Data.Bifunctor (first)
  import Data.ByteString (ByteString)
  import qualified Data.ByteString as B
  import Data.Coerce
  import Data.Default
  import qualified Data.HashMap.Strict as HashMap
  import Data.Hashable (hash)
  import Data.Int
  import Data.IntMap (IntMap)
  import qualified Data.IntMap as IntMap
  import Data.Map (Map)
  import qualified Data.Map as Map
  import Data.Maybe
  import Data.Set (Set)
  import qualified Data.Set as Set
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.Encoding as Text
  import Compat.Prettyprinter hiding ((<>))
  import Compat.Prettyprinter.Render.Text
  import qualified Data.Vector as Vector
  import Data.Vector (Vector)
  import qualified Data.Vector.Storable as VS
  import Data.Word (Word64)
  import System.IO
  import TextShow
  
  import ServiceData.GlobalStats as Stats
  import qualified ServiceData.Types as Stats
  import Util.AllocLimit
  import Util.Timing
  import Util.Log
  import Util.STM
  
  import Glean.Display
  import qualified Glean.Angle.Parser as Angle
  import Glean.Angle.Types hiding (Type, FieldDef, SourcePat_(..))
  import qualified Glean.Angle.Types as Angle
  import qualified Glean.Internal.Types as Thrift
  import qualified Glean.Backend.Types as Backend
  import Glean.Schema.Types
    (RefTarget(..), LookupResult(..), NameEnv, resolveRef)
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Schema.Types
  import Glean.Database.Open
  import qualified Glean.Database.PredicateStats as PredicateStats
  import Glean.Database.Types as Database
  import Glean.Database.Writes
  import Glean.FFI
  import Glean.Query.Codegen
  import Glean.Query.Codegen.Types hiding (Match(..))
  import Glean.Query.Transform
  import Glean.Query.Flatten
  import Glean.Query.Opt
  import Glean.Query.Reorder
  import Glean.Query.Incremental (makeIncremental)
  import Glean.RTS as RTS
  import Glean.RTS.Bytecode.Disassemble
  import qualified Glean.RTS.Bytecode.Gen.Version as Bytecode
  import qualified Glean.RTS.Foreign.Bytecode as Bytecode
  import Glean.RTS.Foreign.FactSet (FactSet)
  import qualified Glean.RTS.Foreign.FactSet as FactSet
  import Glean.RTS.Foreign.Lookup
  import Glean.RTS.Foreign.Ownership
  import Glean.RTS.Foreign.Query
  import Glean.RTS.Foreign.Stacked (stacked)
  import Glean.RTS.Types (Type, PidRef(..))
  import qualified Glean.Types as Thrift
  import qualified Glean.ServerConfig.Types as ServerConfig
  import Glean.Query.JSON
  import Glean.Schema.Resolve
  import Glean.Schema.Util
  import Glean.Util.Observed as Observed
  import Glean.Query.Typecheck
  import Glean.Bytecode.SysCalls (userQuerySysCalls)
  import Glean.Types (UserQueryCont)
  import qualified Glean.Backend.Types as StackedDbOpts
  import Glean.Util.Some
  import Glean.Query.Typecheck.Types (TcPred)
  -- NOTE: We keep the public interface monomorphic, at least for now.
  
  --
  -- | Perform a query using the JSON query syntax, and return the results
  -- as JSON-encoded Thrift data.
  --
  userQuery
    :: Database.Env
    -> Thrift.Repo
    -> Thrift.UserQuery
    -> IO Thrift.UserQueryResults
  userQuery env repo query = chooseEncoding
    (Thrift.userQuery_options query)
    (Thrift.userQuery_encodings query)
    $ genericUserQuery env repo query
  
  --
  -- | Query a list of facts and return the result as JSON-encoded Thrift data.
  --
  userQueryFacts
    :: Database.Env
    -> Thrift.Repo
    -> Thrift.UserQueryFacts
    -> IO Thrift.UserQueryResults
  userQueryFacts env repo query = chooseEncoding
    (Thrift.userQueryFacts_options query)
    (Thrift.userQueryFacts_encodings query)
    $ genericUserQueryFacts env repo query
  
  -- | A generic implementation of scan queries.
  genericUserQuery
    :: Encoding e
    => Database.Env
    -> Thrift.Repo
    -> Thrift.UserQuery
    -> e
    -> IO Thrift.UserQueryResults
  {-# INLINE genericUserQuery #-}
  genericUserQuery env repo query enc = do
    config@ServerConfig.Config{..} <- Observed.get (envServerConfig env)
    readDatabaseWithBoundaries env repo $ \odb bounds lookup ->
      maybe id limitAllocsThrow config_query_alloc_limit
        $ performUserQuery enc (odbSchema odb)
        $ userQueryImpl env odb config NoExtraSteps bounds lookup repo query
  
  -- | A generic implementation of lookup queries.
  genericUserQueryFacts
    :: Encoding e
    => Database.Env
    -> Thrift.Repo
    -> Thrift.UserQueryFacts
    -> e
    -> IO Thrift.UserQueryResults
  {-# INLINE genericUserQueryFacts #-}
  genericUserQueryFacts env repo query enc = do
    config <- Observed.get (envServerConfig env)
    readDatabase env repo $ \odb lookup ->
      performUserQuery enc (odbSchema odb) $
        userQueryFactsImpl (odbSchema odb) config lookup query
  
  -- | Results returned from a query, parametrised of the types of facts and
  -- statistics
  data Results stats fact = Results
    { -- | Facts matching the query
      resFacts :: [(Fid, fact)]
  
      -- | Type of the results, if known
    , resPredicate :: Maybe PredicateDetails
  
      -- | Additional facts satisfying nested parts of the query
    , resNestedFacts :: IntMap fact
  
      -- | Continuation for paging
    , resCont :: Maybe Thrift.UserQueryCont
  
      -- | Statistics
    , resStats :: stats
  
      -- | Diagnostics
    , resDiags :: [Text]
  
      -- | Write handle for derived facts
    , resWriteHandle :: Maybe Thrift.Handle
  
      -- | Count of the facts searched per Pid
    , resFactsSearched :: Maybe (Map Int64 Int64)
  
      -- | Inferred type of the query, for logging
    , resType :: Maybe Text
  
      -- | Size of compiled bytecode, for logging
    , resBytecodeSize :: Maybe Int
  
      -- | Query compile time, for logging, in seconds
    , resCompileTime :: Maybe Double
  
    , resCodegenTime :: Maybe Double
  
      -- | Query execution time after compilation, for logging, in nanoseconds
    , resExecutionTime :: Maybe Word64
  
    -- | Size of query results, for logging
    , resResultBytes :: Maybe Int64
    }
  
  class Encoding e where
    type EncodedFact e
  
    shouldExpand :: e -> Bool
    serializeFact
      :: e
      -> IntMap (EncodedFact e)
      -> Fid
      -> PredicateDetails
      -> Thrift.Fact
      -> IO (EncodedFact e)
    setResults
      :: e
      -> Results stats (EncodedFact e)
      -> Thrift.UserQueryResults
      -> Thrift.UserQueryResults
  
  withEncoding
    :: Thrift.UserQueryEncoding
    -> (forall e. Encoding e => e -> a)
    -> Maybe a
  {-# INLINE withEncoding #-}
  withEncoding (Thrift.UserQueryEncoding_bin x) f = Just $ f x
  withEncoding (Thrift.UserQueryEncoding_json x) f = Just $ f x
  withEncoding (Thrift.UserQueryEncoding_compact x) f = Just $ f x
  withEncoding (Thrift.UserQueryEncoding_listbin x) f = Just $ f x
  withEncoding _ _ = Nothing
  
  chooseEncoding
    :: Maybe Thrift.UserQueryOptions
    -> [Thrift.UserQueryEncoding]
    -> (forall e. Encoding e => e -> IO a)
    -> IO a
  {-# INLINE chooseEncoding #-}
  chooseEncoding mopts encs f = do
    forM_ (Thrift.userQueryOptions_continuation =<< mopts) checkUserQueryCont
    case encs of
      [] -> f $ LegacyJSONEncoding $ fromMaybe def mopts
      _ | act : _ <- mapMaybe (`withEncoding` f) encs -> act
        | otherwise -> throwIO $ Thrift.BadQuery
            "none of the requested encodings are supported"
  
  newtype LegacyJSONEncoding = LegacyJSONEncoding Thrift.UserQueryOptions
  
  instance Encoding LegacyJSONEncoding where
    type EncodedFact LegacyJSONEncoding = ByteString
  
    shouldExpand (LegacyJSONEncoding opts) =
      Thrift.userQueryOptions_expand_results opts
  
    serializeFact
      (LegacyJSONEncoding opts)
      serialized
      fid
      predicateDetails
      Thrift.Fact{..} = factToJSON
        (Thrift.userQueryOptions_no_base64_binary opts)
        serialized
        predicateDetails
        fid
        fact_key
        fact_value
  
    setResults _ res qres = qres
        { Thrift.userQueryResults_facts = snd <$> resFacts res
        , Thrift.userQueryResults_nestedFacts = Map.fromList
            [ (fromIntegral i, fact)
              | (i,fact) <- IntMap.toList $ resNestedFacts res ]
        }
  
  instance Encoding Thrift.UserQueryEncodingJSON where
    type EncodedFact Thrift.UserQueryEncodingJSON = ByteString
  
    shouldExpand = Thrift.userQueryEncodingJSON_expand_results
  
    serializeFact enc serialized fid predicateDetails Thrift.Fact{..} =
      factToJSON
        (Thrift.userQueryEncodingJSON_no_base64_binary enc)
        serialized
        predicateDetails
        fid
        fact_key
        fact_value
  
    setResults enc res qres = qres
      { Thrift.userQueryResults_results = Thrift.UserQueryEncodedResults_json
          Thrift.UserQueryResultsJSON
            { userQueryResultsJSON_encoding = enc
            , userQueryResultsJSON_facts = snd <$> resFacts res
            , userQueryResultsJSON_nestedFacts = Map.fromList
                [ (fromIntegral i, fact)
                  | (i,fact) <- IntMap.toList $ resNestedFacts res ]
            }
      }
  
  instance Encoding Thrift.UserQueryEncodingCompact where
    type EncodedFact Thrift.UserQueryEncodingCompact = ByteString
  
    shouldExpand = Thrift.userQueryEncodingCompact_expand_results
  
    serializeFact _ serialized fid predicateDetails Thrift.Fact{..} =
      factToCompact
        serialized
        predicateDetails
        fid
        fact_key
        fact_value
  
    setResults enc res qres = qres
      { Thrift.userQueryResults_results = Thrift.UserQueryEncodedResults_compact
          Thrift.UserQueryResultsCompact
            { userQueryResultsCompact_encoding = enc
            , userQueryResultsCompact_facts = snd <$> resFacts res
            , userQueryResultsCompact_nestedFacts = Map.fromList
                [ (fromIntegral i, fact)
                  | (i,fact) <- IntMap.toList $ resNestedFacts res ]
            }
      }
  
  instance Encoding Thrift.UserQueryEncodingBin where
    type EncodedFact Thrift.UserQueryEncodingBin = Thrift.Fact
  
    shouldExpand _ = False
  
    serializeFact _ _ _ _ = return
  
    setResults enc res qres = qres
      { Thrift.userQueryResults_results = Thrift.UserQueryEncodedResults_bin
          Thrift.UserQueryResultsBin
            { userQueryResultsBin_encoding = enc
            , userQueryResultsBin_facts =
                Map.fromList
                $ coerce
                $ resFacts res
            , userQueryResultsBin_nestedFacts =
                Map.fromList
                $ map (first fromIntegral)
                $ IntMap.toList
                $ resNestedFacts res
            }
      }
  
  instance Encoding Thrift.UserQueryEncodingListBin where
    type EncodedFact Thrift.UserQueryEncodingListBin = Thrift.Fact
  
    shouldExpand _ = False
  
    serializeFact _ _ _ _ = return
  
    setResults enc res qres = qres
      { Thrift.userQueryResults_results = Thrift.UserQueryEncodedResults_listbin
          Thrift.UserQueryResultsListBin
            { userQueryResultsListBin_encoding = enc
            , userQueryResultsListBin_ids =
                coerce $ Vector.fromList (map fst (resFacts res))
            , userQueryResultsListBin_facts =
                Vector.fromList (map snd (resFacts res))
            , userQueryResultsListBin_nestedFacts =
                Map.fromList
                $ map (first fromIntegral)
                $ IntMap.toList
                $ resNestedFacts res
            }
      }
  
  -- | Perform a query and convert the 'Results' to a Thrift result
  performUserQuery
    :: Encoding e
    => e
    -> DbSchema
    -> IO (Results Stats Thrift.Fact)
    -> IO Thrift.UserQueryResults
  {-# INLINE performUserQuery #-}
  performUserQuery encoding schema query = do
    !results <- withStats $ do
      res <- query
  
      -- Convert nested facts - we do it in the order of their fact ids
      -- which means we can only depend on facts converted earlier.
      nested <- foldM
        (\expanded (fid, fact@Thrift.Fact{..}) -> do
            details <- pidDetails schema fact_type
            encoded <- serializeFact
              encoding
              (if shouldExpand encoding then expanded else IntMap.empty)
              (Fid $ fromIntegral fid)
              details
              fact
            return $ IntMap.insert fid encoded expanded)
        IntMap.empty
        (IntMap.toList $ resNestedFacts res)
  
      let
        -- avoid lookupPid if we know the result type
        getDetails = case resPredicate res of
          Just d -> const (return d)
          Nothing -> pidDetails schema
  
      facts <- forM (resFacts res) $ \(i, fact) -> do
        details <- getDetails (Thrift.fact_type fact)
        enc <- serializeFact
          encoding
          (if shouldExpand encoding then nested else IntMap.empty)
          i
          details
          fact
        return (i,enc)
  
      return res
        { resFacts = facts
        , resNestedFacts = if shouldExpand encoding then IntMap.empty else nested
        }
    return $ setResults encoding results def
      { Thrift.userQueryResults_stats = Just $ resStats results
      , Thrift.userQueryResults_continuation = resCont results
      , Thrift.userQueryResults_diagnostics = resDiags results
      , Thrift.userQueryResults_handle = resWriteHandle results
      , Thrift.userQueryResults_type = resType results
      }
  
  pidDetails :: DbSchema -> Int64 -> IO PredicateDetails
  pidDetails schema ty = do
    case lookupPid (Pid (fromIntegral ty)) schema of
      Nothing ->  throwIO $ Thrift.Exception $
        "failed looking up fact type: " <> Text.pack (show ty)
      Just d -> return d
  
  userQueryFactsImpl
    :: DbSchema
    -> ServerConfig.Config
    -> Lookup
    -> Thrift.UserQueryFacts
    -> IO (Results Stats Thrift.Fact)
       -- The length of the result list is guaranteed to be the same
       -- as the userQueryFacts_facts list in the input.
  userQueryFactsImpl
      schema@DbSchema{..}
      config
      lookup
      query@Thrift.UserQueryFacts{..} = do
    let opts = fromMaybe def userQueryFacts_options
  
    schemaSelector <- schemaVersionForQuery schema config userQueryFacts_schema_id
  
    expandPids <- optsExpandPids opts schemaSelector schema
    let limits = mkQueryRuntimeOptions opts config expandPids
  
    trans <- transformationsForQuery schema schemaSelector
  
    vlog 2 $ "userQueryFactsImpl: " <> show (length userQueryFacts_facts)
    (qResults@QueryResults{..}, fullScans) <- do
      nextId <- firstFreeId lookup
      -- executeCompiled needs a Define, even though we won't use it
      bracket (FactSet.new nextId) release $ \derived -> do
      let stack = stacked lookup derived
      bracket
        (compileQueryFacts userQueryFacts_facts)
        (release . compiledQuerySub) $ \sub -> do
          results <- executeCompiled schemaInventory Nothing stack sub limits
          appliedTrans <- either (throwIO . Thrift.BadQuery) return $
            userQueryFactsTransformations trans schemaSelector schema query results
          -- use Pids in result facts to apply a suitable transformation if neded.
          return
            ( transformResultsBack appliedTrans results
            , compiledQueryFullScans sub
            )
  
    stats <- getStats schema fullScans qResults
  
    let results = Results
          { resFacts = Vector.toList queryResultsFacts
          , resPredicate = Nothing
          , resNestedFacts = mkNestedFacts queryResultsNestedFacts
          , resCont = Nothing
          , resStats = stats
          , resDiags = []
          , resWriteHandle = Nothing
          , resFactsSearched = Nothing
          , resType = Nothing  -- could be facts of different predicates
          , resBytecodeSize = Nothing
          , resCompileTime = Nothing
          , resCodegenTime = Nothing
          , resExecutionTime = Nothing
          , resResultBytes = Nothing
          }
  
    return $ if Thrift.userQueryOptions_omit_results opts
       then withoutFacts results
       else results
  
  userQueryFactsTransformations
    :: QueryTransformations
    -> SchemaSelector
    -> DbSchema
    -> Thrift.UserQueryFacts
    -> QueryResults
    -> Either Text ResultTransformations
  userQueryFactsTransformations qtrans selector schema query results = do
    nameEnv <-  maybe (Left "invalid schema_id") return $ do
      schemaNameEnv schema selector
  
    let predRefs :: [PidRef]
        predRefs = mapMaybe (toPidRef nameEnv) sourceRefs
  
        -- a type containing all types in the response
        allTypes :: Type
        allTypes = Angle.RecordTy
          [ Angle.FieldDef "" (PredicateTy () predId) | predId <- predRefs ]
  
    -- errors if multiple versions of the same predicate are requested.
    transformationsFor schema qtrans allTypes
    where
      sourceRefs :: [SourceRef]
      sourceRefs = mapMaybe toRef $ Set.toList vset
        where
          vset :: Set (Pid, Maybe Version)
          vset = Vector.foldr insert mempty zipped
          zipped = Vector.zip (queryResultsFacts results) versions
          insert ((_, fact), version) acc =
            Set.insert (Pid $ Thrift.fact_type fact, version) acc
  
          versions :: Vector (Maybe Version)
          versions = Vector.fromList $
            Thrift.factQuery_predicate_version <$>
              Thrift.userQueryFacts_facts query
  
          toRef :: (Pid, Maybe Version) -> Maybe SourceRef
          toRef (pid, mversion) = do
            details <- lookupPid pid schema
            let PredicateRef name _ = predicateRef details
            return $ SourceRef name mversion
  
      toPidRef :: NameEnv RefTargetId -> SourceRef -> Maybe PidRef
      toPidRef nameEnv ref = do
        ResolvesTo (RefPred predId) <- return $ resolveRef nameEnv ref
        details <- lookupPredicateId predId schema
        return (pidRef details)
  
  -- | A version of userQuery where we only care about the resulting writes
  -- caused by the query. Used for stored predicate derivation.
  userQueryWrites
    :: Database.Env
    -> OpenDB s
    -> ServerConfig.Config
    -> Boundaries
    -> Lookup
    -> Thrift.Repo
    -> PredicateId
    -> Thrift.UserQuery
    -> IO (
         Thrift.UserQueryStats,
         Maybe Thrift.UserQueryCont,
         Maybe Thrift.Handle
       )
  userQueryWrites env odb config bounds lookup repo pred q = do
    meta <- atomically $ Catalog.readMeta (envCatalog env) repo
    inc <- shouldDeriveIncrementally meta
    mode <-
      if inc
      then IncrementalDerivation <$> sectionsStats
      else return NoExtraSteps
    Results{..} <- withStats $
      userQueryImpl env odb config mode bounds lookup repo q
    return (resStats, resCont, resWriteHandle)
    where
      -- We derive incrementally if
      --   1. the DB is stacked, and
      --   2. the predicate was derived in the base DB
      --
      -- (2) might be false if we're deriving a predicate in a
      -- stacked DB that we didn't derive on the base DB.
      shouldDeriveIncrementally meta =
        case Thrift.metaDependencies meta of
          Just (Thrift.Dependencies_stacked Thrift.Stacked{..}) ->
            check $ Thrift.Repo stacked_name stacked_hash
          Just (Thrift.Dependencies_pruned pruned) ->
            check (Thrift.pruned_base pruned)
          Nothing -> return False
        where
        check base = do
          meta <- atomically $ Catalog.readMeta (envCatalog env) base
          return $ predicateIdRef pred `elem` Thrift.metaCompletePredicates meta
  
      -- Here the best we can do is say whether a Pid could have facts in a DB.
      -- Because of ownership slicing it may be that even though
      -- `Backend.predicateStats` gives a positive number all facts could be
      -- filtered out.
      sectionsStats :: IO (SeekSection -> Pid -> Bool)
      sectionsStats = do
        stackStats <- stats Backend.IncludeBase
        topStats <- stats Backend.ExcludeBase
        let baseStats = stackStats `minus` topStats
            pidHasFacts seekWhere (Pid pid) = case seekWhere  of
              SeekOnAllFacts -> hasFacts pid stackStats
              SeekOnBase -> hasFacts pid baseStats
              SeekOnStacked -> hasFacts pid topStats
        return pidHasFacts
  
      stats opts =
        fmap Thrift.predicateStats_count
        <$> PredicateStats.predicateStats env repo opts
      minus = Map.unionWith (-)
      hasFacts pid m = maybe False (> 0) $ Map.lookup pid m
  
  userQueryImpl
    :: Database.Env
    -> OpenDB s
    -> ServerConfig.Config
    -> CompilationMode
    -> Boundaries
    -> Lookup
    -> Thrift.Repo
    -> Thrift.UserQuery
    -> IO (Results Stats Thrift.Fact)
  
  userQueryImpl
    env
    odb
    config
    mode
    bounds
    lookup
    repo
    query@Thrift.UserQuery{..} = do
      let opts = fromMaybe def userQuery_options
  
      case Thrift.userQueryOptions_syntax opts of
        Thrift.QuerySyntax_ANGLE -> return ()
        other -> throwIO $ Thrift.Exception $
          "query syntax not supported: " <> Text.pack (show other)
  
      let
        schema = odbSchema odb
        stored = Thrift.userQueryOptions_store_derived_facts opts
        debug = Thrift.userQueryOptions_debug opts
  
      schemaVersion <-
        schemaVersionForQuery schema config userQuery_schema_id
  
      (returnType,compileTime,diag,cont) <-
        case Thrift.userQueryOptions_continuation opts of
          Just ucont
            | Just retTy <- Thrift.userQueryCont_returnType ucont -> do
            (compileTime, _, returnType) <-
              timeIt $ compileType schema schemaVersion retTy
            return (returnType,compileTime,[],Right ucont)
  
          -- This is either a new query or the continuation of a query
          -- that returns a temporary predicate.
          _ -> do
            (compileTime, _, (query@QueryWithInfo{..}, ty, preds)) <-
              timeIt $ compileAngleQuery
                (envEnableRecursion env)
                schemaVersion
                schema
                mode
                userQuery_query
                stored
                (envDebug env)
  
            predDiag <- if Thrift.queryDebugOptions_pred_has_facts debug then
                getPredDiags env repo schema preds
                else return []
            let
              irDiag =
                [ "ir:\n" <> Text.pack (show (displayDefault qiQuery))
                | Thrift.queryDebugOptions_ir debug ]
  
              cont = case Thrift.userQueryOptions_continuation opts of
                Just c -> Right c
                Nothing -> Left query
            return (ty,compileTime,irDiag <> predDiag,cont)
  
      details <- getReturnPredicateDetails schema returnType
  
      if Thrift.userQueryOptions_just_check opts then do
        return $ emptyResult {resDiags = diag}
      else do
        let compileInfo = CompileInfo {
          returnType = returnType,
          compileTime = compileTime,
          compileDiag = diag,
          cont = cont
          }
        runQuery env odb config bounds lookup repo details compileInfo query
  
  data CompileInfo = CompileInfo {
    returnType :: Type,
    compileTime :: Double,
    compileDiag :: [Text],
    cont :: Either CodegenQuery UserQueryCont
  }
  
  getPredDiags
    :: Database.Env
    -> Thrift.Repo
    -> DbSchema
    -> [TcPred]
    -> IO [Text]
  getPredDiags env repo schema preds = do
    predStats <- PredicateStats.predicateStats env repo StackedDbOpts.IncludeBase
    let predsMissing = filter (\p -> not (tcPredExist p predStats)) preds
    return $ map warnDiag predsMissing
      where
        tcPredExist (PidRef pid _, _) predStats =
          derived schema pid || hasFacts predStats pid
        -- we assume derived predicates "exist" since we won't find any facts
        derived schema pid = case lookupPid pid schema of
          Just details -> case predicateDeriving details of
            Derive DeriveOnDemand _ -> True
            Derive DeriveIfEmpty _ -> True
            _ -> False
          Nothing -> error "Unknown Pid in getPredDiags"
        hasFacts predStats pid = case Map.lookup (fromPid pid) predStats of
          Just stat -> Thrift.predicateStats_count stat > 0
          Nothing -> False
        warnDiag :: TcPred -> Text
        warnDiag (PidRef _ (PredicateId predRef _), Some span) = Text.pack $
          "Warning: "
          <> show (pretty span) ++ " "
          <> Text.unpack (predicateRef_name predRef)
          <> Text.unpack " doesn't exist in " ++ show (Thrift.repo_name repo)
  
  getReturnPredicateDetails :: DbSchema -> Type -> IO PredicateDetails
  getReturnPredicateDetails schema@DbSchema{..} returnType = do
    case returnType of
      Angle.PredicateTy _ (PidRef pid _) ->
        case IntMap.lookup (fromIntegral (fromPid pid)) predicatesByPid of
          Nothing -> throwIO $ Thrift.Exception "internal: no predicate"
          Just d -> return d
  
      _not_a_predicate -> do
        return
            -- This is a temporary Pid generated by Flatten.captureKey,
            -- because the query doesn't return facts. In this case we
            -- are pretending the query returns facts for a new predicate,
            -- so we need a PredicateDetails to pass to unfoldFacts below.
            -- (this is temporary, unfoldFacts wil go away in due course
            -- when we compile it to bytecode).
            PredicateDetails
              { predicatePid = tempPid schema
              , predicateKeyType = returnType
              , predicateValueType = unit
              , predicateId = tempPredicateId
              , predicateSchema = error "predicateSchema"
              , predicateTraversal = error "predicateTraversal"
              , predicateTypecheck = error "predicateTypecheck"
              , predicateDeriving = NoDeriving
              , predicateInStoredSchema = False
              }
  
  runQuery
    :: Database.Env
    -> OpenDB s
    -> ServerConfig.Config
    -> Boundaries
    -> Lookup
    -> Thrift.Repo
    -> PredicateDetails
    -> CompileInfo
    -> Thrift.UserQuery
    -> IO (Results Stats Thrift.Fact)
  runQuery
    env
    odb
    config
    bounds
    lookup
    repo
    details
    CompileInfo{..}
    Thrift.UserQuery{..} = do
      vlog 2 $ "return type: " <> show (displayDefault returnType)
  
      let
        schema@DbSchema{..} = odbSchema odb
        opts = fromMaybe def userQuery_options
        stored = Thrift.userQueryOptions_store_derived_facts opts
  
      schemaVersion <-
          schemaVersionForQuery schema config userQuery_schema_id
      trans <- transformationsForQuery schema schemaVersion
  
      unless (Text.null userQuery_predicate) $ do
        -- With Angle queries, setting userQuery_predicate is
        -- optional. If the client sets it, we will check it: this
        -- can be a useful way to catch errors in the client.
        -- If the query is not returning whole facts, then the
        -- client should set this field to "".
        let ref = SourceRef userQuery_predicate userQuery_predicate_version
        checkPredicatesMatch schema details ref schemaVersion
  
      expandPids <- optsExpandPids opts schemaVersion schema
      let limits = mkQueryRuntimeOptions opts config expandPids
  
      nextId <- case Thrift.userQueryOptions_continuation opts of
        Just Thrift.UserQueryCont{..}
          | userQueryCont_nextId > 0 -> return (Fid userQueryCont_nextId)
        _otherwise -> firstFreeId lookup
      derived <- FactSet.new nextId
      let stack = stacked lookup derived
  
      defineOwners <- if stored
        then do
          maybeOwnership <- readTVarIO (odbOwnership odb)
          forM maybeOwnership $ \ownership ->
            newDefineOwnership ownership nextId
        else return Nothing
  
      appliedTrans <- either (throwIO . Thrift.BadQuery) return $
        transformationsFor schema trans returnType
  
      ( qResults@QueryResults{..}
        , queryDiag
        , bytecodeSize
        , codegenTime
        , fullScans ) <-
        case cont of
          Right ucont -> do
            let binaryCont = Thrift.userQueryCont_continuation ucont
            results <- transformResultsBack appliedTrans <$>
              restartCompiled
                schemaInventory
                defineOwners
                stack
                (Just $ predicatePid details)
                limits
                binaryCont
            return (results, [], B.length binaryCont, 0, [])
  
          Left query -> do
            let
              debug = Thrift.userQueryOptions_debug opts
              bytecodeDiag sub =
                [ "bytecode:\n" <> Text.unlines
                  (disassemble "Query" userQuerySysCalls $ compiledQuerySub sub)
                | Thrift.queryDebugOptions_bytecode debug ]
  
            bracket
              (timeIt $ compileQuery (envEnableRecursion env) trans bounds query)
              (\(_, _, sub) -> release $ compiledQuerySub sub)
              $ \(codegenTime, _, sub) -> do
                results <- transformResultsBack appliedTrans <$>
                  executeCompiled schemaInventory defineOwners stack sub limits
  
                diags <-
                  evaluate $ force (bytecodeDiag sub) -- don't keep sub alive
                sz <- evaluate $ Bytecode.size (compiledQuerySub sub)
                let fullScans = compiledQueryFullScans sub
                return (results, diags, sz, codegenTime, fullScans)
  
      -- If we're storing derived facts, queue them for writing and
      -- return the handle. We allow querying for stored derived
      -- predicates with stored=True on a read-only DB; this is used
      -- by the regression testing framework to test derived predicates.
      maybeWriteHandle <-
        if stored && (
            isJust (odbWriting odb) ||
            Thrift.userQueryOptions_omit_results opts)
          then writeDerivedFacts env repo nextId derived defineOwners
            queryResultsFacts
          else return Nothing
  
      userCont <- case queryResultsCont of
        Nothing -> return Nothing
        Just bs -> do
          nextId <- firstFreeId derived
          return $ Just $ mkUserQueryCont (Right returnType) bs nextId
  
      stats <- getStats schema fullScans qResults
  
      when (isJust userCont) $
        addStatValueType "glean.query.truncated" 1 Stats.Sum
  
      let ppType = renderStrict $ layoutPretty defaultLayoutOptions $
            displayDefault returnType
  
          results = Results
            { resFacts = Vector.toList queryResultsFacts
            , resPredicate = Just details
            , resNestedFacts = mkNestedFacts queryResultsNestedFacts
            , resCont = userCont
            , resStats = stats
            , resDiags = compileDiag ++ queryDiag
            , resWriteHandle = maybeWriteHandle
            , resFactsSearched = queryResultsStats
            , resType = Just ppType
            , resBytecodeSize = Just bytecodeSize
            , resCompileTime = Just compileTime
            , resCodegenTime = Just codegenTime
            , resExecutionTime = Just queryResultsElapsedNs
            , resResultBytes = Just queryResultsResultBytes
            }
  
      return $ if Thrift.userQueryOptions_omit_results opts
        then withoutFacts results
        else results
  
  transformationsForQuery
    :: DbSchema
    -> SchemaSelector
    -> IO QueryTransformations
  transformationsForQuery schema selector = do
    case allSchemaVersion schema selector of
      Nothing -> throwIO $ Thrift.BadQuery "invalid schema_id"
      Just schemaId ->
        let transMap = predicatesTransformations schema in
        case HashMap.lookup schemaId transMap of
          Just trans -> return trans
          Nothing -> throwIO $ Thrift.BadQuery "no transformations for schema_id"
  
  schemaVersionForQuery
    :: DbSchema
    -> ServerConfig.Config
    -> Maybe Thrift.SchemaId  -- ^ SchemaId specified by client
    -> IO SchemaSelector
  schemaVersionForQuery schema ServerConfig.Config{..} qid = do
    use <-
      case qid of
        Nothing -> return LatestSchema
        Just id
          | id `Map.member` schemaEnvs schema -> return (SpecificSchemaId id)
          | config_strict_query_schema_id ->
              throwIO (Thrift.UnknownSchemaId id)
          | otherwise -> do
              logWarning $ "schema unavailable: " <> show id
              return LatestSchema
  
    vlog 1 $ "using schema ID: " <> show (pretty use)
    return use
  
  optsExpandPids
    :: Thrift.UserQueryOptions
    -> SchemaSelector
    -> DbSchema
    -> IO (Set Pid)
  optsExpandPids opts schemaVersion dbSchema =
    fmap Set.fromList $ forM (Thrift.userQueryOptions_expand_predicates opts) $
      \(Thrift.SourcePredicate name version) -> do
        let ref = SourceRef name version
        case lookupPredicateSourceRef ref schemaVersion dbSchema of
          Left err -> throwIO $ Thrift.BadQuery err
          Right details -> return (predicatePid details)
  
  data CompilationMode
    = NoExtraSteps
    | IncrementalDerivation (SeekSection -> Pid -> Bool)
  
  compileAngleQuery
    :: EnableRecursion
    -> SchemaSelector
      -- ^ Schema version to resolve unversioned predicates
    -> DbSchema
    -> CompilationMode
      -- ^ only used in predicate derivations on incremental dbs
    -> ByteString
    -> Bool
    -> DebugFlags
    -> IO (CodegenQuery, Type, [TcPred])
  compileAngleQuery rec ver dbSchema mode source stored debug = do
    parsed <- checkBadQuery Text.pack $ Angle.parseQuery source
    ifDebug $ "parsed query: " <> show (displayDefault parsed)
  
    let scope = addTmpPredicate $ fromMaybe HashMap.empty $
          schemaNameEnv dbSchema ver
  
    resolved <- checkBadQuery id $ runExcept $
      runResolve latestAngleVersion scope $ resolveQuery parsed
    ifDebug $ "resolved query: " <> show (displayDefault resolved)
  
    (typechecked, preds) <- (checkBadQuery id =<<) $ runExceptT $
      typecheck dbSchema (defaultTcOpts debug latestAngleVersion)
        (dbSchemaRtsType dbSchema) resolved
    ifDebug $ "typechecked query: " <> show (displayDefault (qiQuery typechecked))
  
    flattened <- checkBadQuery id $ runExcept $
      flatten rec dbSchema latestAngleVersion stored typechecked
    ifDebug $ "flattened query: " <> show (displayDefault (qiQuery flattened))
  
    optimised <- checkBadQuery id $ runExcept $ optimise flattened
    ifDebug $ "optimised query: " <> show (displayDefault (qiQuery optimised))
  
    reordered <- checkBadQuery id $ runExcept $ reorder dbSchema optimised
    ifDebug $ "reordered query: " <> show (displayDefault (qiQuery reordered))
  
    final <- case mode of
      NoExtraSteps -> return reordered
      IncrementalDerivation getStats -> do
        vlog 2 "made incremental"
        return $ makeIncremental getStats reordered
  
    return (final, qiReturnType typechecked, preds)
    where
    ifDebug = when (queryDebug debug) . hPutStrLn stderr
  
    checkBadQuery :: (err -> Text) -> Either err a -> IO a
    checkBadQuery txt act = case act of
      Left str -> throwIO $ Thrift.BadQuery $ txt str
      Right a -> return a
  
  -- | Put the nested facts in the right form for the conversion to
  -- (JSON, Compact, Bin).
  mkNestedFacts :: Vector (Fid, Thrift.Fact) -> IntMap Thrift.Fact
  mkNestedFacts facts =
    IntMap.fromList
      [ (fromIntegral (fromFid id), f)
      | (id,f) <- Vector.toList facts ]
  
  withoutFacts :: Results stats fact -> Results stats fact
  withoutFacts results = results
      { resFacts = mempty
      , resNestedFacts = mempty
      }
  
  mkQueryRuntimeOptions
    :: Thrift.UserQueryOptions
    -> ServerConfig.Config
    -> Set Pid
    -> QueryRuntimeOptions
  mkQueryRuntimeOptions
      Thrift.UserQueryOptions{..} ServerConfig.Config{..} expandPids =
    QueryRuntimeOptions
      { queryMaxResults = userQueryOptions_max_results
          <|> config_default_max_results -- from ServerConfig
      , queryMaxBytes = userQueryOptions_max_bytes
          <|> config_default_max_bytes -- from ServerConfig
      , queryMaxTimeMs = userQueryOptions_max_time_ms
          <|> config_default_max_time_ms -- from ServerConfig
      , queryMaxSetSize = config_max_set_size_bytes
      , queryWantStats = userQueryOptions_collect_facts_searched
      , queryDepth = if
          | userQueryOptions_recursive && not userQueryOptions_omit_results ->
            ExpandRecursive
          | not (null expandPids) -> ExpandPartial expandPids
          | otherwise -> ResultsOnly
      }
  
  emptyResult :: Results Stats fact
  emptyResult = Results {
      resFacts = mempty
    , resPredicate = Nothing
    , resNestedFacts = mempty
    , resCont = Nothing
    , resStats = Stats {
        statFactCount = 0
      , statResultCount = 0
      , statFullScans = []
      }
    , resDiags = []
    , resWriteHandle = Nothing
    , resFactsSearched = Nothing
    , resType = Nothing
    , resBytecodeSize = Nothing
    , resCompileTime = Nothing
    , resCodegenTime = Nothing
    , resExecutionTime = Nothing
    , resResultBytes = Nothing
    }
  
  
  {- Note [Writing derived facts]
  
  When deriving, we
    (1) Run a query that produces all the derived facts
    (2) Serialize the FactSet
    (3) Write this to the DB
  
  As the query runs, it writes derived facts into a FactSet and also
  produces a set of results. The results are correct, however the
  FactSet may contain some additional invalid facts. For example, this
  can happen if the derived predicate is something like
  
  predicate P : string
    X where
      Q X;
      X != "foo"
  
  when we compile the query, it can look something like
  
  F where
    Q X;
    F = P<- X;    -- (*)
    X != "foo"
  
  The statement labelled (*) produces the derived fact. Note that this
  is just a regular statement; it can get arbitrarily reordered relative
  to the other statements. In this case the compiler chose to put it
  before the final filter X != "foo", which means that we'll create some
  derived facts that aren't actually true. There's nothing wrong with
  this: the result of the query is still correct, but we have to be
  careful to use the query results and not just write the whole contents
  of the FactSet to the DB.
  
  It would be difficult to ensure the compiler always ordered F = P<- X
  after the filter, there's nothing requiring it to do this. We can't
  use statement ordering for this, because statement ordering is for
  optimisation and shouldn't change sementics. The compiler can override
  ordering hints if it wants.
  
  The right thing to do is to use the query results to filter the
  FactSet, and fortunately there's a good place to do this: we already
  reorder the results when we serialize the FactSet (for ownership
  reasons), so we can do the filtering there.
  -}
  
  writeDerivedFacts
    :: Env
    -> Thrift.Repo
    -> Fid
    -> FactSet
    -> Maybe DefineOwnership
    -> Vector (Fid, Thrift.Fact)
    -> IO (Maybe Thrift.Handle)
  writeDerivedFacts env repo firstId derived owned results = do
    -- See Note [Writing Derived Facts] for why we filter here
    let order = VS.fromList
           [ fromFid fid
           | (fid,_) <- Vector.toList results
           , fid >= firstId  -- ignore existing facts
           ]
    batch <- case owned of
      Nothing -> FactSet.serializeReorder derived order
      Just define -> do
        nextId <- firstFreeId derived
        let count = fromIntegral (fromFid nextId - fromFid firstId)
        sorted <- defineOwnershipSortByOwner define count order
        FactSet.serializeReorder derived sorted
    -- If the batch is empty, we may still have new ownership data about
    -- existing facts, so we have to write that
    if Thrift.batch_count batch == 0 && isNothing owned
      then return Nothing
      else do
        resp <- enqueueBatch env Thrift.ComputedBatch
          { computedBatch_repo = repo
          , computedBatch_remember = True
          , computedBatch_batch = batch }
          owned
        case resp of
          Thrift.SendResponse_handle handle -> return (Just handle)
          Thrift.SendResponse_retry (Thrift.BatchRetry s) ->
            throwIO $ Thrift.Retry s
  
  
  -- | Check that the predicate declared in the UserQuery request
  -- matches the actual type of the query as determined by the
  -- typechecker. This adds some rudimentary type safety: the predicate
  -- declared in the UserQuery is typically the one expected at the call
  -- site.
  checkPredicatesMatch
    :: DbSchema
    -> PredicateDetails
    -> SourceRef
    -> SchemaSelector
    -> IO ()
  checkPredicatesMatch dbSchema details predicate schemaVer = do
    case lookupPredicateSourceRef predicate schemaVer dbSchema of
      Left err -> throwIO $ Thrift.BadQuery err
      Right details' ->
        if predicatePid details == predicatePid details'
          then return ()
          else throwIO $ Thrift.BadQuery $ mconcat
            [ "predicate " <> showRef predicate <>
              " does not match type of query: "
            <> predicateRef_name (predicateRef details) <> "."
            <> showt (predicateRef_version (predicateRef details)) ]
  
  
  data Stats = Stats
    { statFactCount :: {-# UNPACK #-} !Int
    , statResultCount :: {-# UNPACK #-} !Int
    , statFullScans :: [PredicateRef]
    }
  
  getStats :: DbSchema -> [Pid] -> QueryResults -> IO Stats
  getStats schema fullScans QueryResults{..} = do
    let
      results =
        Vector.length queryResultsFacts
  
      facts =
        Vector.length queryResultsFacts +
        Vector.length queryResultsNestedFacts
  
      pref pid = case lookupPid pid schema of
        Nothing -> error "Unknown Pid in getStats"
        Just details -> predicateIdRef $ predicateId details
  
    addStatValueType "glean.query.facts" facts Stats.Sum
    addStatValueType "glean.query.results" results Stats.Sum
    return $ Stats
      { statFactCount = facts
      , statResultCount  = results
      , statFullScans = map pref fullScans
      }
  
  withStats :: IO (Results Stats fact) -> IO (Results Thrift.UserQueryStats fact)
  withStats io = do
    (secs, bytes, res) <- timeIt io
    let stats = Thrift.UserQueryStats
          { Thrift.userQueryStats_num_facts =
              fromIntegral $ statFactCount $ resStats res
          , Thrift.userQueryStats_result_count =
              fromIntegral $ statResultCount $ resStats res
          , Thrift.userQueryStats_elapsed_ns = truncate (secs * 1000000000)
          , Thrift.userQueryStats_allocated_bytes = fromIntegral bytes
          , Thrift.userQueryStats_facts_searched = resFactsSearched res
          , Thrift.userQueryStats_bytecode_size =
              fromIntegral <$> resBytecodeSize res
          , Thrift.userQueryStats_compile_time_ns =
              fmap (round . (* 1000000000)) (resCompileTime res )
          , Thrift.userQueryStats_codegen_time_ns =
              fmap (round . (* 1000000000)) (resCodegenTime res )
          , Thrift.userQueryStats_execute_time_ns =
              fromIntegral <$> resExecutionTime res
          , Thrift.userQueryStats_full_scans = statFullScans $ resStats res
          , Thrift.userQueryStats_result_bytes =
              fromIntegral <$> resResultBytes res
          }
    return res{ resStats = stats }
  
  mkUserQueryCont
    :: Either (Set Pid) Type
    -> ByteString
    -> Fid
    -> Thrift.UserQueryCont
  mkUserQueryCont contInfo cont nextId =
    hashUserQueryCont $ Thrift.UserQueryCont
    { userQueryCont_continuation = cont
    , userQueryCont_nextId = fromFid nextId
    , userQueryCont_version = fromIntegral Bytecode.version
    , userQueryCont_hash = 0
    , userQueryCont_returnType = returnType
    , userQueryCont_pids = pids
    }
    where
      returnType = case contInfo of
        Right ty -> Just $ serializeType ty
        Left _ -> Nothing
      pids = case contInfo of
        Right _ -> []
        Left ps -> map fromPid $ Set.elems ps
  
  checkUserQueryCont :: Thrift.UserQueryCont -> IO ()
  checkUserQueryCont cont@Thrift.UserQueryCont{..} = do
    when
      (userQueryCont_version < fromIntegral Bytecode.lowestSupportedVersion
        || userQueryCont_version > fromIntegral Bytecode.version)
      $ throwIO $ Thrift.BadQuery $
          "unsupported query continuation version "
            <> Text.pack (show userQueryCont_version)
    when
      (userQueryCont_hash /= Thrift.userQueryCont_hash (hashUserQueryCont cont))
      $ throwIO $ Thrift.BadQuery "invalid query continuation hash"
  
  hashUserQueryCont :: Thrift.UserQueryCont -> Thrift.UserQueryCont
  hashUserQueryCont Thrift.UserQueryCont{..} = Thrift.UserQueryCont
    -- NOTE: The hash is really just a checksum to detect accidental corruption
    -- so a 64 bit non-crypto hash should work fine. Hashable currently uses a
    -- slightly broken (https://github.com/tibbe/hashable/issues/190) version of
    -- FNV-1 but should still be good enough.
    { userQueryCont_hash = fromIntegral $ hash
        ( userQueryCont_continuation
        , userQueryCont_nextId
        , userQueryCont_version
        , userQueryCont_returnType
        )
    , ..
    }
  
  serializeType :: Type -> ByteString
  serializeType = Text.encodeUtf8 . renderStrict . layoutCompact .
    displayDefault
      -- Note that we have to print PredicateIds as PredicateRefs, because
      -- the parser will parse it as a source type. The default display options
      -- do exactly that.
  
  compileType :: DbSchema -> SchemaSelector -> ByteString -> IO Type
  compileType schema version src = do
    parsed <- checkParsed $ Angle.parseType src
    let scope = addTmpPredicate $ fromMaybe HashMap.empty $
          schemaNameEnv schema version
    resolved <- checkResolved $ runResolve latestAngleVersion scope $
      resolveType parsed
    checkConverted $ dbSchemaRtsType schema resolved
    where
      checkParsed = either (badQuery . Text.pack) return
      checkResolved = either badQuery return . runExcept
      checkConverted = maybe (badQuery "type not present in schema") return
  
      badQuery :: Text -> IO a
      badQuery err = throwIO $ Thrift.BadQuery $ Text.unlines
        ["unable to compile type: ", err, "type was: ", Text.decodeUtf8 src]
  