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.Prune (pruneDerivations) where
  
  import Control.Monad
  import Control.Monad.State (State, runState)
  import qualified Control.Monad.State as State
  import Data.Bitraversable (bitraverse)
  import Data.Foldable (asum)
  import Data.Maybe (fromMaybe)
  import qualified Data.Set as Set
  import qualified Data.Graph as Graph
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import qualified Data.IntMap.Strict as IntMap
  import Data.IntMap.Strict (IntMap)
  
  import Glean.Database.Schema.Types (PredicateDetails(..))
  import Glean.RTS.Types (PidRef(..), Type)
  import Glean.Angle.Types (PredicateId(..), DerivingInfo(..))
  import qualified Glean.Angle.Types as Angle
  import Glean.Query.Codegen.Types
  import Glean.Query.Typecheck (tcQueryDeps)
  import Glean.Query.Typecheck.Types
    (TcPat, TcTerm(..), TcStatement(..), TcQuery(..), TypecheckedQuery)
  import Glean.RTS.Term (Term(..))
  
  -- | Remove all branches that we know will not succeed given the content of the
  -- database.
  pruneDerivations
    :: (PredicateId -> Bool)
    -> HashMap PredicateId PredicateDetails
    -> HashMap PredicateId PredicateDetails
  pruneDerivations hasStoredFacts details =
    updateDerivation <$> details
    where
      updateDerivation details = fromMaybe details $ do
        Derive when _ <- return $ predicateDeriving details
        let newDerivation =
              case HashMap.lookup (predicateId details) pruned of
                Nothing -> NoDeriving
                Just query -> Derive when query
        return $ details { predicateDeriving = newDerivation }
  
      pruned :: HashMap PredicateId TypecheckedQuery
      pruned = pruneQueries hasStoredFacts derivations
  
      derivations :: HashMap PredicateId TypecheckedQuery
      derivations = flip HashMap.mapMaybe details $ \PredicateDetails{..} ->
        case predicateDeriving of
          NoDeriving -> Nothing
          Derive _ query -> Just query
  
  pruneQueries
    :: (PredicateId -> Bool)
    -> HashMap PredicateId TypecheckedQuery
    -> HashMap PredicateId TypecheckedQuery
  pruneQueries hasStoredFacts derivations =
    foldr add mempty $ topologicalSort derivations
    where
      add :: (PredicateId, TypecheckedQuery)
          -> HashMap PredicateId TypecheckedQuery
          -> HashMap PredicateId TypecheckedQuery
      add (predId, query) derivedWithFacts =
         case prune (hasFacts predId derivedWithFacts) query of
            Nothing -> derivedWithFacts
            Just pruned -> HashMap.insert predId pruned derivedWithFacts
  
      hasFacts predId derivedWithFacts child =
        child == predId
          || hasStoredFacts child
          || child `HashMap.member` derivedWithFacts
  
  topologicalSort
    :: HashMap PredicateId TypecheckedQuery
    -> [(PredicateId, TypecheckedQuery)]
  topologicalSort derivations =
    [ entry
    | vertex <- Graph.topSort graph
    , let (entry,_,_) = fromVertex vertex
    ]
    where
      (graph, fromVertex, _) = Graph.graphFromEdges edges
      edges =
        [ (entry, predId, dependencies)
        | entry@(predId, query) <- HashMap.toList derivations
        , let dependencies = Set.toList $ tcQueryDeps $ qiQuery query
        ]
  
  -- | Remove paths that we know will not yield any result.
  -- This will save us a lot of compilation at query time.
  --
  -- We can determine that a path is doomed to fail if it has a FactGenerator
  -- that will search on a predicate with no facts.
  --
  -- Pruning won't leave us with unbound variables because:
  -- - If a statement in a sequence is set to fail, the entire sequence is removed
  -- - If one sequence is removed from a disjunction, all variables bound by it
  --   should also be bound by the other disjunction alternatives
  prune :: (PredicateId -> Bool) -> TypecheckedQuery -> Maybe TypecheckedQuery
  prune hasFacts (QueryWithInfo q _ gen t) = do
    renumberVars gen t <$> pruneTcQuery q
    where
    pruneTcQuery :: TcQuery -> Maybe TcQuery
    pruneTcQuery (TcQuery ty keyPat mvalPat stmts ord) =
      TcQuery ty
        <$> prunePat keyPat
        <*> case mvalPat of
              Nothing -> Just Nothing
              Just v -> Just <$> prunePat v
        <*> pruneSequence stmts
        <*> pure ord
  
    pruneSequence :: [TcStatement] -> Maybe [TcStatement]
    pruneSequence = traverse pruneStmt
  
    pruneStmt :: TcStatement -> Maybe TcStatement
    pruneStmt (TcStatement ty lhs rhs) =
      TcStatement ty <$> prunePat lhs <*> prunePat rhs
  
    prunePat :: TcPat -> Maybe TcPat
    prunePat pat = case pat of
      Byte{} -> Just pat
      Nat{} -> Just pat
      ByteArray{} -> Just pat
      String{} -> Just pat
      Array xs -> Array <$> traverse prunePat xs
      Tuple xs -> Tuple <$> traverse prunePat xs
      Alt i x -> Alt i <$> prunePat x
      Ref m -> case m of
        MatchWild{} -> Just pat
        MatchNever{} -> Nothing
        MatchFid{} -> Just pat
        MatchBind{} -> Just pat
        MatchVar{} -> Just pat
        MatchAnd a b -> Ref <$> (MatchAnd <$> prunePat a <*> prunePat b)
        MatchPrefix s x -> Ref . MatchPrefix s <$> prunePat x
        MatchArrayPrefix t xs all ->
          Ref <$> (MatchArrayPrefix t <$> traverse prunePat xs <*> prunePat all)
        MatchExt (Typed ty tcterm) -> case tcterm of
          TcFactGen (PidRef _ predId) _ _ _
            | not $ hasFacts predId
            -> Nothing
          TcFactGen pidref k v range -> Ref . MatchExt . Typed ty
            <$> (TcFactGen pidref <$> prunePat k <*> prunePat v <*> pure range)
          TcElementsOfArray x -> Ref . MatchExt . Typed ty . TcElementsOfArray
            <$> prunePat x
          TcElementsOfSet x ->
            Ref . MatchExt . Typed ty . TcElementsOfSet <$> prunePat x
          TcElementsUnresolved ty' x ->
            Ref . MatchExt . Typed ty . TcElementsUnresolved ty' <$> prunePat x
          TcQueryGen q ->
            Ref . MatchExt . Typed ty . TcQueryGen <$> pruneTcQuery q
          -- we dont' want to handle negation here because if it tries to match
          -- against things that are not in the database it should succeed.
          TcAll query ->
            Ref . MatchExt . Typed ty . TcAll <$> pruneTcQuery query
          TcNegation{} -> Just pat
          TcPrimCall op xs -> Ref . MatchExt . Typed ty . TcPrimCall op
            <$> traverse prunePat xs
          TcIf (Typed t' c) t e ->
            let cond = prunePat c
                condt = Typed t' <$> cond
                then_ = prunePat t
                else_ = prunePat e
            in
            asum
              [ Ref . MatchExt . Typed ty <$> (TcIf <$> condt <*> then_ <*> else_)
              , Ref <$> (MatchAnd <$> cond <*> then_)
              , else_ :: Maybe TcPat
              ]
          TcOr a b ->
            let pa = prunePat a
                pb = prunePat b
            in
            asum
              [ Ref . MatchExt . Typed ty <$> (TcOr <$> pa <*> pb)
              , pa
              , pb
              ]
          TcDeref ty' p
            | Angle.PredicateTy _ (PidRef _ predId) <-
              ty', not $ hasFacts predId -> Nothing
            | otherwise ->
              Ref . MatchExt . Typed ty . TcDeref ty' <$> prunePat p
          TcFieldSelect (Typed ty' p) f -> do
            p' <- prunePat p
            return $ Ref $ MatchExt $ Typed ty $ TcFieldSelect (Typed ty' p') f
          TcAltSelect (Typed ty' p) f -> do
            p' <- prunePat p
            return $ Ref $ MatchExt $ Typed ty $ TcAltSelect (Typed ty' p') f
          TcPromote _ p -> prunePat p
          TcDemote _ p -> prunePat p
          TcStructPat{} -> error "prune: TcStructPat"
  
  type R a = State S a
  
  data S = S
    { nextVar :: Int
    , mappings :: IntMap Int
    }
  
  -- | After removing branches from the query we must now update
  -- variable names to ensure we use the smallest numbers possible.
  renumberVars :: Maybe Generator -> Type -> TcQuery -> TypecheckedQuery
  renumberVars gen ty q =
    let (newQuery, S varCount _) = runState (renameQuery q) (S 0 mempty)
    in
    QueryWithInfo newQuery varCount gen ty
    where
    renameQuery :: TcQuery -> R TcQuery
    renameQuery (TcQuery ty key mval stmts ord) =
      TcQuery ty
        <$> renamePat key
        <*> traverse renamePat mval
        <*> traverse renameStmt stmts
        <*> pure ord
  
    renameStmt :: TcStatement -> R TcStatement
    renameStmt (TcStatement ty lhs rhs) =
      TcStatement ty <$> renamePat lhs <*> renamePat rhs
  
    renamePat :: TcPat -> R TcPat
    renamePat = traverse (bitraverse renameTyped renameVar)
  
    renameTyped :: Typed TcTerm -> R (Typed TcTerm)
    renameTyped (Typed ty term) = Typed ty <$> renameTcTerm term
  
    renameTcTerm :: TcTerm -> R TcTerm
    renameTcTerm = \case
      TcOr a b -> TcOr <$> renamePat a <*> renamePat b
      TcFactGen ref k v range ->
        TcFactGen ref <$> renamePat k <*> renamePat v <*> pure range
      TcElementsOfArray x -> TcElementsOfArray <$> renamePat x
      TcElementsOfSet x -> TcElementsOfSet <$> renamePat x
      TcElementsUnresolved ty x -> TcElementsUnresolved ty <$> renamePat x
      TcQueryGen q -> TcQueryGen <$> renameQuery q
      TcAll query -> TcAll <$> renameQuery query
      TcNegation xs -> TcNegation <$> traverse renameStmt xs
      TcPrimCall op xs -> TcPrimCall op <$> traverse renamePat xs
      TcIf cond then_ else_ ->
        TcIf <$> traverse renamePat cond <*> renamePat then_ <*> renamePat else_
      TcDeref ty p -> TcDeref ty <$> renamePat p
      TcFieldSelect (Typed ty p) f -> do
        p' <- renamePat p
        return $ TcFieldSelect (Typed ty p') f
      TcAltSelect (Typed ty p) f -> do
        p' <- renamePat p
        return $ TcAltSelect (Typed ty p') f
      TcPromote ty p -> TcPromote ty <$> renamePat p
      TcDemote ty p -> TcDemote ty <$> renamePat p
      TcStructPat fs -> fmap TcStructPat $ forM fs $ \(n,p) ->
        (n,) <$> renamePat p
  
    renameVar :: Var -> R Var
    renameVar (Var ty old n) = State.state $ \s ->
      case IntMap.lookup old (mappings s) of
        Just new -> (Var ty new n, s)
        Nothing ->
          let new = nextVar s
              mappings' = IntMap.insert old new (mappings s)
              next' = new + 1
          in
          (Var ty new n, S next' mappings')
  