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.Expand
    ( expandDerivedPredicateCall
    ) where
  
  import Control.Monad.State
  import Data.Bifunctor (first)
  
  import Glean.Query.Codegen.Types
    (Match(..), Var(..), Pat, QueryWithInfo(..), Typed(..))
  import Glean.Query.Flatten.Types
  import Glean.Query.Typecheck.Types
  import Glean.Query.Vars
  import Glean.RTS.Term as RTS
  import Glean.Database.Schema.Types
  
  
  expandDerivedPredicateCall
    :: PredicateDetails
    -> Pat                -- ^ pattern at the call site
    -> Pat                -- ^ value pattern, if applicable
    -> TypecheckedQuery   -- ^ query from the derived predicate
    -> F TcQuery
  expandDerivedPredicateCall PredicateDetails{..} key val QueryWithInfo{..} = do
    (TcQuery _ keyDef maybeValDef stmts ord) <-
      instantiateWithFreshVariables qiQuery qiNumVars
  
    let
      key' = fmap (first (\_ -> error "MatchExt")) key
      val' = fmap (first (\_ -> error "MatchExt")) val
  
    -- we have
    --    pred P -> Q
    -- where
    --    pred : T
    --       K -> V where S
    --
    -- let's turn it into
    --
    --   { X, Y } where S; X = K, X = P; Y = V; Y = Q;
    --
    -- If the predicate has no value, we can simplify
    --
    --   X where S; X = K, X = P
  
    case maybeValDef of
      Just valDef -> do
        x <- fresh predicateKeyType
        y <- fresh predicateValueType
        return $
          TcQuery predicateKeyType (Ref (MatchVar x)) (Just (Ref (MatchVar y)))
            (stmts ++ [
              TcStatement predicateKeyType (Ref (MatchBind x)) keyDef,
              TcStatement predicateKeyType (Ref (MatchVar x)) key',
              TcStatement predicateValueType (Ref (MatchBind y)) valDef,
              TcStatement predicateValueType (Ref (MatchVar y)) val' ])
            ord
      Nothing -> do
        x <- fresh predicateKeyType
        return $
          TcQuery predicateKeyType (Ref (MatchVar x)) Nothing
            (stmts ++ [
              TcStatement predicateKeyType (Ref (MatchBind x)) keyDef,
              TcStatement predicateKeyType (Ref (MatchVar x)) key' ])
            ord
  
  -- | Make a fresh instance of a query where none of the variables
  -- clash with existing variables. We know the maximum variable in the
  -- query, so we just reserve this number of variables and add the
  -- current value of exNextVar to each variable in the query.
  instantiateWithFreshVariables :: TcQuery -> Int -> F TcQuery
  instantiateWithFreshVariables query numVars = do
    state <- get
    let base = flNextVar state
    put state { flNextVar = base + numVars }
    return $ instantiateQuery base query
    where
    instantiateQuery base (TcQuery ty head maybeVal stmts ord) =
      TcQuery ty
        (instantiatePat base head)
        (fmap (instantiatePat base) maybeVal)
        (map (instantiateStmt base) stmts)
        ord
  
    instantiateStmt base (TcStatement ty lhs rhs) =
      TcStatement ty (instantiatePat base lhs) (instantiatePat base rhs)
  
    instantiateTcTerm base (TcOr a b) =
      TcOr (instantiatePat base a) (instantiatePat base b)
    instantiateTcTerm base (TcFactGen pid kpat vpat range) =
      TcFactGen pid (instantiatePat base kpat) (instantiatePat base vpat) range
    instantiateTcTerm base (TcElementsOfArray pat) =
      TcElementsOfArray (instantiatePat base pat)
    instantiateTcTerm base (TcElementsOfSet pat) =
      TcElementsOfSet (instantiatePat base pat)
    instantiateTcTerm base (TcElementsUnresolved ty pat) =
      TcElementsUnresolved ty (instantiatePat base pat)
    instantiateTcTerm base (TcQueryGen query) =
      TcQueryGen (instantiateQuery base query)
    instantiateTcTerm base (TcAll query) =
      TcAll (instantiateQuery base query)
    instantiateTcTerm base (TcNegation stmts) =
      TcNegation (map (instantiateStmt base) stmts)
    instantiateTcTerm base (TcPrimCall op args) =
      TcPrimCall op (map (instantiatePat base) args)
    instantiateTcTerm base (TcIf (Typed ty cond) then_ else_) =
      TcIf
        (Typed ty $ instantiatePat base cond)
        (instantiatePat base then_)
        (instantiatePat base else_)
    instantiateTcTerm base (TcDeref ty pat) =
      TcDeref ty (instantiatePat base pat)
    instantiateTcTerm base (TcFieldSelect (Typed ty pat) field) =
      TcFieldSelect (Typed ty (instantiatePat base pat)) field
    instantiateTcTerm base (TcAltSelect (Typed ty pat) field) =
      TcAltSelect (Typed ty (instantiatePat base pat)) field
    instantiateTcTerm base (TcPromote ty pat) =
      TcPromote ty (instantiatePat base pat)
    instantiateTcTerm base (TcDemote ty pat) =
      TcDemote ty (instantiatePat base pat)
    instantiateTcTerm base (TcStructPat fs) =
      TcStructPat [(n, instantiatePat base p) | (n,p) <- fs]
  
    instantiatePat :: Int -> TcPat -> TcPat
    instantiatePat base = fmap (instantiateMatch base)
  
    instantiateMatch base m = case m of
      MatchVar (Var ty v nm) -> MatchVar (Var ty (v+base) nm)
      MatchBind (Var ty v nm) -> MatchBind (Var ty (v+base) nm)
      MatchAnd l r -> MatchAnd (instantiatePat base l) (instantiatePat base r)
      MatchPrefix s t -> MatchPrefix s (instantiatePat base t)
      MatchExt (Typed ty ext) -> MatchExt (Typed ty (instantiateTcTerm base ext))
      MatchArrayPrefix ty pre all -> MatchArrayPrefix ty (instantiatePat base <$> pre) (instantiatePat base all)
      other -> other
  