{-
  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.Flatten
  ( flatten
  ) where

import Control.Monad
import Control.Monad.Except
import Control.Monad.State
import Data.List hiding (intersect)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe
import Data.Text (Text)
import qualified Data.Text as Text

import qualified Glean.Angle.Types as Angle
import Glean.Display
import Glean.Query.Codegen.Types
  ( Match(..)
  , Var(..)
  , Pat
  , Generator_(..)
  , Generator
  , Expr
  , Typed(..)
  , SeekSection(..)
  , QueryWithInfo(..))
import Glean.Query.Expand
import Glean.Query.Flatten.Types
import Glean.Query.Typecheck.Types
import Glean.Query.Vars
import Glean.RTS.Types as RTS
import Glean.RTS.Term as RTS
import Glean.Database.Schema.Types
import Glean.Database.Types (EnableRecursion(..))
import qualified Glean.Angle.Types as Schema
import Glean.Schema.Util

-- | Turn 'TypecheckedQuery' into 'FlattenedQuery', by lifting out
-- nested generators into statements.
flatten
  :: EnableRecursion
  -> DbSchema
  -> Schema.AngleVersion
  -> Bool -- ^ derive DerivedAndStored predicates
  -> TypecheckedQuery
  -> Except Text FlattenedQuery
flatten rec dbSchema ver deriveStored QueryWithInfo{..} =
  fmap fst $ flip runStateT state $ do
    (flattened, maybeLookup, returnType) <- do
      flat <- flattenQuery qiQuery `catchError` flattenFailure
      captureKey ver dbSchema flat (case qiQuery of TcQuery ty _ _ _ _ -> ty)
    nextVar <- gets flNextVar
    return $ QueryWithInfo flattened nextVar maybeLookup returnType
  where
      state = initialFlattenState rec dbSchema qiNumVars deriveStoredPred

      deriveStoredPred =
        case derefType qiReturnType of
          Angle.PredicateTy _ (PidRef _ pref) | deriveStored -> Just pref
          _ -> Nothing

      flattenFailure e = throwError $
        e <> " in\n" <> Text.pack (show (displayVerbose qiQuery))

flattenQuery :: TcQuery -> F FlatQuery
flattenQuery query = do
  (group, head, maybeVal) <- flattenQuery' query
  return (FlatQuery head maybeVal group)

flattenQuery' :: TcQuery -> F (FlatStatementGroup, Expr, Maybe Expr)
flattenQuery' (TcQuery ty head Nothing stmts ord) = do
  (ords, floats) <- mapAndUnzipM flattenStatement stmts
  pats <- flattenPattern head
  case pats of
    [(stmts,head')] -> return $ case ord of
      Angle.Ordered -> (mkGroup ords (stmts : floats), head', Nothing)
      Angle.Unordered -> (mkGroup [] (stmts : ords <> floats), head', Nothing)
    _many -> do -- TODO: ord
      -- If there are or-patterns on the LHS, then we have
      --    P1 | P2 | ... where stmts
      -- so we will generate
      --    X where stmts; X = P1 | P2 | ...
      v <- fresh ty
      let
        alts =
          [ asGroup (stmts `thenStmt`
              FlatStatement ty (Ref (MatchBind v)) (TermGenerator head))
          | (stmts, head) <- pats ]
      return
        ( case ord of
            Angle.Ordered ->
              mkGroup ords (oneStmt (FlatDisjunction alts) : floats)
            Angle.Unordered ->
              mkStatementGroup
                (Floating (FlatDisjunction alts) :
                  flattenStmts (mconcat (ords <> floats)))
        , Ref (MatchVar v)
        , Nothing
        )
flattenQuery' (TcQuery ty head (Just val) stmts _ord {- TODO -}) = do
  (ords, floats) <- mapAndUnzipM flattenStatement stmts
  pats <- flattenPattern head
  vals <- flattenPattern val
  case [ (stmtsp <> stmtsv, head, val)
       | (stmtsp, head) <- pats
       , (stmtsv, val) <- vals ] of
    [(stmts, head', val)] ->
      return (mkGroup ords (stmts : floats), head', Just val)
    many -> do
      -- As above, but we must handle the value too
      k <- fresh ty
      v <- fresh ty
      let
        alts =
          [ asGroup (stmts `thenStmt`
              FlatStatement ty lhsPair (TermGenerator (Tuple [head,val])))
          | (stmts, head, val) <- many ]
        lhsPair = Tuple [Ref (MatchBind k), Ref (MatchBind v)]
      return
        ( mkGroup ords (oneStmt (FlatDisjunction alts) : floats)
        , Ref (MatchVar k)
        , Just (Ref (MatchVar v))
        )

flattenStatement :: TcStatement -> F (Statements, Statements)
flattenStatement (TcStatement ty lhs rhs) = do
  rgens <- flattenSeqGenerators rhs
  lgens <- flattenSeqGenerators lhs
  mkStmt ty lgens rgens

-- | Here we decide whether an explicit statement is ordered or
-- unordered with respect to other statements in the group. Implicit
-- statements (e.g. those resulting from nested patterns) are always
-- unordered.
addStmt
  :: FlatStatement
  -> (Statements, Statements)
  -> (Statements, Statements)
addStmt stmt@(FlatStatement _ _ TermGenerator{}) (here, float) =
  (here, float `thenStmt` stmt)  -- filters are always unordered
addStmt stmt (here, float) =
  (here `thenStmt` stmt, float)

mkStmt
  :: Type
  -> [(Statements, Statements, Generator)]
  -> [(Statements, Statements, Generator)]
  -> F (Statements, Statements)
mkStmt ty [(lhshere, lhsfloat, TermGenerator lhs)]
    [(rhshere, rhsfloat, gen)] = do
  return $ addStmt (FlatStatement ty lhs gen)
    (lhshere <> rhshere, lhsfloat <> rhsfloat)
mkStmt ty [(lhshere, lhsfloat, gen)]
    [(rhshere, rhsfloat, TermGenerator rhs)] = do
  return $ addStmt (FlatStatement ty rhs gen)
    (lhshere <> rhshere, lhsfloat <> rhsfloat)
mkStmt ty [(lhshere, lhsfloat, lgen)] [(rhshere, rhsfloat, rgen)] = do
  v <- fresh ty
  return $
    addStmt (FlatStatement ty (Ref (MatchBind v)) lgen) $
    addStmt (FlatStatement ty (Ref (MatchBind v)) rgen)
      (lhshere <> rhshere, lhsfloat <> rhsfloat)
mkStmt ty [(lhshere, lhsfloat, gen)] many = do
  v <- fresh ty
  return $
    addStmt (FlatStatement ty (Ref (MatchBind v)) gen)
      (lhshere
       `thenStmt` disjunction
            [ mkGroup
                [rhshere]
                [rhsfloat `thenStmt`
                  FlatStatement ty (Ref (MatchBind v)) rhsgen ]
            | (rhshere, rhsfloat, rhsgen) <- many
            ]
      , lhsfloat
      )
mkStmt ty many [one] = mkStmt ty [one] many
mkStmt ty lhsmany rhsmany = do
  v <- fresh ty
  (lhsord, lhsfloat) <- bind v lhsmany
  (rhsord, rhsfloat) <- bind v rhsmany
  return (lhsord <> rhsord, lhsfloat <> rhsfloat)

flattenSeqGenerators :: TcPat -> F [(Statements, Statements, Generator)]
flattenSeqGenerators (Ref (MatchExt (Typed ty match))) = case match of
  TcOr left right -> do
    l <- flattenSeqGenerators left
    r <- flattenSeqGenerators right
    return (l ++ r)
  TcFactGen pid kpat vpat range -> do
    kpats <- flattenPattern kpat
    vpats <- flattenPattern vpat
    sequence
      [ do
          (stmts, gen) <- flattenFactGen pid range kpat vpat
          return (stmts, kstmts <> vstmts, gen)
      | (kstmts, kpat) <- kpats
      , (vstmts, vpat) <- vpats ]
  TcElementsOfArray pat -> do
    r <- flattenPattern pat
    return [(mempty, stmts, ArrayElementGenerator ty pat') | (stmts,pat') <- r ]
  TcElementsOfSet pat -> do
    r <- flattenPattern pat
    return [(mempty, stmts, SetElementGenerator ty pat') | (stmts,pat') <- r ]
  TcQueryGen query -> do
    (group, term, _) <- flattenQuery' query
    return [(floatGroup group, mempty, TermGenerator term)]
  TcAll query -> do
    (group, term, _) <- flattenQuery' query
    var <- fresh ty
    return
      [ (Statements [FlatAllStatement var term group]
      , mempty
      , TermGenerator (Ref (MatchVar var)))]
  TcNegation stmts -> do
    (ords, floats) <- mapAndUnzipM flattenStatement stmts
    let neg = FlatNegation (mkGroup ords floats)
    return [(oneStmt neg, mempty, TermGenerator $ Tuple [])]
  TcPrimCall op args -> do
    r <- manyTerms (\args -> PrimCall op args ty) <$> mapM flattenPattern args
    return [ (mempty, float, gen) | (float, gen) <- r ]
  TcIf (Typed condTy cond) then_ else_ -> do
    (condOrd, condFloat) <- bindWild condTy =<< flattenSeqGenerators cond
    var <- fresh ty
    (thenOrd, thenFloat) <- bind var =<< flattenSeqGenerators then_
    (elseOrd, elseFloat) <- bind var =<< flattenSeqGenerators else_
    let stmt = FlatConditional
          (asGroup (condOrd <> condFloat))
          (asGroup (thenOrd <> thenFloat))
          (asGroup (elseOrd <> elseFloat))
    return [(mempty `thenStmt` stmt, mempty,
      TermGenerator $ Ref $ MatchVar var)]
  _other -> do
    r <- flattenPattern (Ref (MatchExt (Typed ty match)))
    return $ [(mempty, stmts, TermGenerator pat) | (stmts, pat) <- r]
flattenSeqGenerators pat = do
  r <- flattenPattern pat
  return $ [(mempty, stmts, TermGenerator pat) | (stmts, pat) <- r]

bindWild
  :: Type
  -> [(Statements, Statements, Generator)]
  -> F (Statements, Statements)
bindWild ty stmts =
  mkStmt ty stmts [(mempty, mempty, TermGenerator $ Ref $ MatchWild ty)]

bind
  :: Var
  -> [(Statements, Statements, Generator)]
  -> F (Statements, Statements)
bind var gens = mkStmt
  (varType var)
  [(mempty, mempty, TermGenerator (RTS.Ref (MatchBind var)))]
  gens

flattenFactGen
  :: PidRef
  -> SeekSection
  -> Pat
  -> Pat
  -> F (Statements, Generator)
flattenFactGen pidRef@(PidRef pid _) rng kpat vpat = do
  dbSchema <- gets flDbSchema
  deriveStored <- gets flDeriveStored
  case lookupPid pid dbSchema of
    Nothing -> lift $ throwError $
      "internal error: flatten: " <> Text.pack (show pid)
    Just details@PredicateDetails{..} -> do
      let factGen = (mempty, FactGenerator pidRef kpat vpat rng)
      case predicateDeriving of
        Schema.NoDeriving ->
          return factGen
        Schema.Derive when query
          | Schema.DerivedAndStored <- when
          , Just predicateId /= deriveStored ->
               return factGen
          | otherwise -> do
            calling predicateId factGen $ do
              query' <- expandDerivedPredicateCall details kpat vpat query
              (group, key, maybeVal) <- flattenQuery' query'
              let val = fromMaybe (Tuple []) maybeVal
              return (floatGroup group, DerivedFactGenerator pidRef key val)

calling
  :: Schema.PredicateId
  -> a   -- ^ use this value if we already expanded a recursive call.
  -> F a
  -> F a
calling ref seek inner = do
  stack <- gets flStack
  recursion <- gets flRecursion
  if
    | ref `notElem` stack -> do
      modify $ \state -> state { flStack = ref : stack }
      a <- inner
      modify $ \state -> state { flStack = stack }
      return a
    | EnableRecursion <- recursion -> return seek
    | otherwise ->
      throwError $ "recursive reference to predicate " <>
        Text.pack (show (displayDefault ref))

-- | Returns a list of statement*pattern pairs
--   representing a disjunction @(P1 where S1 | P2 where S2 | ...)@
--   where none of @Pi@ nor @Si@ contain any nested generators.
--
--   Moreover, if @Si@ succeeds, it will bind variables in @Pi@.
flattenPattern :: TcPat -> F [(Statements, Pat)]
flattenPattern pat = case pat of
  Byte n -> singleTerm (Byte n)
  Nat n -> singleTerm (Nat n)
  String s -> singleTerm (String s)
  ByteArray s -> singleTerm (ByteArray s)
  Array terms -> manyTerms Array <$> mapM flattenPattern terms
  Tuple terms -> manyTerms Tuple <$> mapM flattenPattern terms
  Alt s x -> fmap (fmap (Alt s)) <$> flattenPattern x
  Ref (MatchWild ty) -> singleTerm (Ref (MatchWild ty))
  Ref (MatchNever ty) -> singleTerm (Ref (MatchNever ty))
  Ref (MatchFid fid) -> singleTerm (Ref (MatchFid fid))
  Ref (MatchBind var) -> singleTerm (Ref (MatchBind var))
  Ref (MatchVar var) -> singleTerm (Ref (MatchVar var))
  Ref (MatchAnd a b) ->
    twoTerms (\a b -> Ref (MatchAnd a b))
      <$> flattenPattern a
      <*> flattenPattern b
  Ref (MatchPrefix str t) ->
    fmap (fmap (Ref . MatchPrefix str)) <$> flattenPattern t
  Ref (MatchArrayPrefix ty prefix all) -> do
    twoTerms (\pre all -> Ref (MatchArrayPrefix ty pre all))
      <$> (manyTerms id <$> mapM flattenPattern prefix)
      <*> flattenPattern all
  Ref (MatchExt (Typed _ (TcOr a b))) -> do  -- Note [flattening TcOr]
    as <- flattenPattern a
    bs <- flattenPattern b
    return (as ++ bs)

  --  *(pat : P) ==> X where pat = P X
  Ref (MatchExt (Typed keyTy (TcDeref ty pat))) -> do
    r <- flattenPattern pat
    ref@(PidRef _ pred)  <- case ty of
      Angle.PredicateTy _ ref -> return ref
      _other -> throwError "TcDeref: not a predicate"
    PredicateDetails{..} <- getPredicateDetails pred
    forM r $ \(stmts, p) -> do
      v <- Ref . MatchVar <$> fresh keyTy
      let valPat = Ref (MatchWild predicateValueType)
          gen = FactGenerator ref v valPat SeekOnAllFacts
      return (stmts `thenStmt` FlatStatement ty p gen, v)

  -- pat.field ==> X where { field = X } = pat
  Ref (MatchExt (Typed ty (TcFieldSelect (Typed recTy pat) name))) -> do
    r <- flattenPattern pat
    case derefType recTy of
      Angle.RecordTy fields -> do
        let sel v =
              [ if name == n then v else Ref (MatchWild ty)
              | Angle.FieldDef n ty <- fields
              ]
        forM r $ \(stmts, p) -> do
          v <- Ref . MatchVar <$> fresh ty
          let stmt = FlatStatement recTy (Tuple (sel v)) (TermGenerator p)
          return (stmts `thenStmt` stmt, v)
      _other ->
        throwError $ "internal: TcFieldSelect: " <>
          Text.pack (show (displayDefault recTy))

  -- pat.field? ==> X where { field = X } = pat
  Ref (MatchExt (Typed ty (TcAltSelect (Typed sumTy pat) name))) -> do
    r <- flattenPattern pat
    n <- case derefType sumTy of
      Angle.SumTy fields
        | (_,n):_ <- lookupField name fields ->
          return n
      _ -> error "flatten: SumTy"
    forM r $ \(stmts, p) -> do
      v <- Ref . MatchVar <$> fresh ty
      let stmt = FlatStatement sumTy (Alt n v) (TermGenerator p)
      return (stmts `thenStmt` stmt, v)

  Ref (MatchExt (Typed _ (TcPromote _ pat))) ->
    flattenPattern pat

  Ref (MatchExt (Typed ty _)) -> do
    gens <- flattenSeqGenerators pat
    v <- fresh ty
    let gen = TermGenerator (RTS.Ref (MatchBind v))
    (ord, float) <- mkStmt ty [(mempty, mempty, gen)] gens
    return [(ord <> float, RTS.Ref (MatchVar v))]

{- Note [flattening TcOr]

To flatten

   p | q

then we have a choice of two plans.

PLAN A

Is just to generate a fresh variable and bind the term. In an expression
context this would mean

   X where X = p | q

(side note: this might not work in a pattern context, e.g. if p/q
contain wildcards or unbound variables, in which case we could instead
do something like

   X where _ = (() where p = X) | (() where q = X)

but the statement must come *after* matching the current pattern,
which we don't currently support. Also note here that we can't easily
tell whether p/q contain unbound variables, because the binding order
at this point is wrong due to substitution and floating, and won't be
correct until we've run fixMatchBind later.)

PLAN B is to duplicate the enclosing pattern, so if we had

   { "a", p | q }

then we'll get

   { "a", p } | { "a", q }

and so on, all the way up to the nearest enclosing statement, where we
can generate SeqGenerators.

It turns out that PLAN B is what we want if we have an enclosing fact
generator, e.g.

   cxx1.Name ("foo".. | "bar"..)

should generate

   (cxx1.Name "foo"..) | (cxx1.Name "bar"..)

because that would be a lot more efficient than

   cxx1.Name X
   _ = (() where "foo".. = X) | (() where "bar".. = X)

However, PLAN A is good in some other cases because it doesn't
duplicate anything.

For now we do PLAN B all the time, but someday we might try to do PLAN
A when it's a good idea.

Note that this only applies to |. The other kinds of things we want
to lift out can always be bound and floated, because they're always
valid expressions and lifting out won't affect performance.
-}

-- | A set of statements. The statements in the set will be reordered
-- by the Reorder pass later.
newtype Statements = Statements { _unStatements :: [FlatStatement] }

instance Semigroup Statements where
  Statements s1 <> Statements s2 = Statements (s2 <> s1)

instance Monoid Statements where
  mempty = Statements []
  mconcat = foldl' (<>) mempty  -- override default, we want left-fold

thenStmt :: Statements -> FlatStatement -> Statements
thenStmt ss s | irrelevant s = ss
thenStmt (Statements ss) s = Statements (s : ss)

oneStmt :: FlatStatement -> Statements
oneStmt s = mempty `thenStmt` s

-- | True for statements that cannot make any difference to the query.
-- That is, statements that can neither fail nor bind any variables.
irrelevant :: FlatStatement -> Bool
irrelevant (FlatStatement _ _ (TermGenerator (Ref (MatchWild _)))) = True
irrelevant (FlatStatement _ (Ref (MatchWild _)) rhs) =
  case rhs of
    DerivedFactGenerator{} -> True
    _ -> False
    -- it's tempting to include PrimCall and TermGenerator here, but
    -- note that these may be written by the user and may be
    -- unresolvable, so elimiating them here may hide an error.
irrelevant _ = False

-- | Inject an ordered sequence of FlatStatementGroup into a
-- Statements.  This is used when we need to retain the ordering
-- between some statements, but allow the whole sequence to be
-- reordered with respect to other statements around it.
floatGroup :: FlatStatementGroup -> Statements
floatGroup g = Statements [grouping g]
  -- Note: we nest groups by using FlatDisjunction with a single
  -- alternative. This is so that this set of groups may be reordered with
  -- respect to other statements/groups at the same level. For this to
  -- work, we have to retain this grouping until the Reorder phase, so
  -- the optimiser must not flatten it away.

flattenStmts :: Statements -> [Ordered FlatStatement]
flattenStmts (Statements s) = map Floating (reverse s)

disjunction :: [FlatStatementGroup] -> FlatStatement
disjunction [FlatStatementGroup [x]] = unOrdered x
disjunction groups = FlatDisjunction groups

mkGroup :: [Statements] -> [Statements] -> FlatStatementGroup
mkGroup ords floats =
  mkStatementGroup $
    flattenStmts (mconcat floats) <>
    map Ordered (orderedGroups ords)
  where
  orderedGroups :: [Statements] -> [FlatStatement]
  orderedGroups = mapMaybe (mkGroup . flattenStmts)

  mkGroup [] = Nothing
  mkGroup [Floating one] = Just one
  mkGroup more  = Just (FlatDisjunction [mkStatementGroup more])

asGroup :: Statements -> FlatStatementGroup
asGroup one = mkStatementGroup (flattenStmts one)

singleTerm :: a -> F [(Statements, a)]
singleTerm t = return [(mempty, t)]

manyTerms :: ([a] -> b) -> [[(Statements, a)]] -> [(Statements, b)]
manyTerms constr subterms =
  [ (mconcat stmts, constr ts)
  | xs <- sequence subterms  -- all combinations of subterms
  , let (stmts,ts) = unzip xs
  ]

twoTerms
  :: (a -> b -> c)
  -> [(Statements, a)]
  -> [(Statements, b)]
  -> [(Statements, c)]
twoTerms constr as bs =
  [ (stmtsa <> stmtsb, constr a b)
  | (stmtsa, a) <- as
  , (stmtsb, b) <- bs
  ]


-- | Note [query result]
--
-- A bit of trickiness to handle capturing and returning the key and
-- value. This avoids having to separately look up the fact later, and
-- also means we can track the size of the output that we're returning
-- in the compiled query.
--
-- For a query like
--
-- > predicate ("x",_)
--
-- flattening expands this to
--
-- > X where X = predicate ("x",_)
--
-- and here we expand this to
--
-- > { X, K, () } where X = predicate K @ ("x",_)
--
-- and the query compiler compiles this to efficient bytecode to
-- capture the key and return it.
--
-- Note that we might not be able to do this transformation if the
-- query isn't in the right form, so in that case we need an
-- explicit fact lookup, like
--
--   X = predicate K
--
-- BUT we have to do this lookup as the very last thing before
-- returning the result, because the fact might be a derived fact
-- produced by a DerivedFactGen, so it won't exist until after it has
-- been produced. Therefore we can't just add this lookup statement to
-- the query, because there's nothing preventing later phases from
-- moving it too early.
--
-- So instead we defer the generation of this statement until code
-- generation. The generator is kept in the field qiGenerator in the
-- QueryWithInfo.

captureKey
  :: Schema.AngleVersion
  -> DbSchema
  -> FlatQuery
  -> Type
  -> F (FlatQuery, Maybe Generator, Type)
captureKey ver dbSchema
    (FlatQuery pat Nothing (FlatStatementGroup ord)) ty
  | Angle.PredicateTy _ pidRef@(PidRef pid _) <- ty  = do
  let
    captureOrdStmt fidVar keyVar maybeValVar (Ordered s) =
      (Ordered <$> stmts, pats)
      where (stmts, pats) = captureStmt fidVar keyVar maybeValVar s
    captureOrdStmt fidVar keyVar maybeValVar (Floating s) =
      (Floating <$> stmts, pats)
      where (stmts, pats) = captureStmt fidVar keyVar maybeValVar s

    -- look for $result = pred pat
    -- replace it with  $result = pred ($key @ pat)
    captureStmt
      :: Int
      -> Var
      -> Maybe Var
      -> FlatStatement
      -> (NonEmpty FlatStatement, Maybe (Pat, Pat))
    captureStmt fidVar keyVar maybeValVar
      (FlatStatement ty lhs (FactGenerator pid kpat vpat range))
      | Ref (MatchBind (Var _ v' _)) <- lhs
      , fidVar == v' =
        let kpat' = Ref (MatchAnd (Ref (MatchBind keyVar)) kpat)
            vpat' = case maybeValVar of
              Nothing -> vpat
              Just valVar -> Ref (MatchAnd (Ref (MatchBind valVar)) vpat)
            stmt' = FlatStatement ty lhs (FactGenerator pid kpat' vpat' range)
            keyExpr = Ref (MatchVar keyVar)
            valExpr = maybe (Tuple []) (Ref . MatchVar) maybeValVar
        in
        (stmt' :| [], Just (keyExpr, valExpr))

    captureStmt fidVar keyVar@(Var keyTy _ _) maybeValVar
      (FlatStatement ty lhs (DerivedFactGenerator pid kexpr vexpr))
      | Ref (MatchBind (Var _ v' _)) <- lhs
      , fidVar == v' =
        let
          -- optimise away a redundant copy when the key/value is
          -- already a variable.
          (keyExpr, bindKey)
            | Ref MatchVar{} <- kexpr = (kexpr, Nothing)
            | otherwise =
              ( Ref (MatchVar keyVar)
              , Just $
                  FlatStatement keyTy
                    (Ref (MatchBind keyVar))
                    (TermGenerator kexpr)
              )
          (valExpr, bindVal)
            | Ref MatchVar{} <- vexpr = (vexpr, Nothing)
            | otherwise = case maybeValVar of
                Nothing -> (Tuple [], Nothing)
                Just valVar@(Var valTy _ _) ->
                  ( Ref (MatchVar valVar)
                  , Just $
                      FlatStatement valTy
                        (Ref (MatchBind valVar))
                        (TermGenerator vexpr)
                  )
        in
        ( NonEmpty.fromList $ catMaybes
            [ bindKey
            , bindVal
            , Just $ FlatStatement ty lhs
                (DerivedFactGenerator pid keyExpr valExpr)
            ]
        , Just (keyExpr, valExpr))
    captureStmt _ _ _ other =
        (other :| [], Nothing)

  PredicateDetails{..} <- case lookupPid pid dbSchema of
    Nothing -> throwError "internal: captureKey"
    Just details -> return details
  keyVar <- fresh predicateKeyType
  maybeValVar <- if eqType ver predicateValueType unit
    then return Nothing
    else Just <$> fresh predicateValueType
  let
    (ord', capturedOrd) =
      case pat of
        Ref (MatchVar (Var _ v _)) ->
          let k :: [Ordered FlatStatement]
                -> [(NonEmpty (Ordered FlatStatement), Maybe (Pat, Pat))]
              k = map (captureOrdStmt v keyVar maybeValVar)
              (ords, capturedOrd) = unzip (k ord)

              conc :: [NonEmpty (Ordered FlatStatement)] ->
                [Ordered FlatStatement]
              conc ((x :| xs) : ys) = x : (xs ++ concatMap NonEmpty.toList ys)
              conc [] = []
          in
          (conc ords, capturedOrd)
        _other -> (ord, [])

    returnTy = tupleSchema [ty, predicateKeyType, predicateValueType]

  case catMaybes capturedOrd of
    [(key, val)] ->
      return (FlatQuery (RTS.Tuple [pat, key, val]) Nothing
        (mkStatementGroup ord'), Nothing, returnTy)
    _ -> do
      -- If we can't find the key/value now, we'll add a statement
      --   X = pred K V
      -- at the end right before codeGen. Here we produce the generator
      -- for the statement.
      let
        gen = FactGenerator pidRef
          (Ref (MatchBind keyVar))
          (Ref (maybe (MatchWild predicateValueType) MatchBind maybeValVar))
          SeekOnAllFacts
      return (FlatQuery pat Nothing (mkStatementGroup ord),
        Just gen,
        returnTy)

  | otherwise = do
  -- We have
  --
  --    E where stmts
  --
  -- we'll transform it into
  --
  --    ($fid, $key) where
  --    stmts;
  --    $key = E
  --    $fid = Pid<- $key
  --
  -- where
  --    Pid is a new Pid for a fictitious derived predicate
  --    $fid is the fact ID
  --    $key is the key
    let
      pidRef = PidRef (tempPid dbSchema) tempPredicateId
      pidTy = Angle.PredicateTy () pidRef
      retTy = tupleSchema [pidTy, ty, unit]

    fidVar <- fresh pidTy
    keyVar <- fresh ty

    let
      result = RTS.Tuple
        [ RTS.Ref (MatchVar fidVar)
        , RTS.Ref (MatchVar keyVar)
        , Tuple []
        ]
      resultStmt1 =
        FlatStatement ty (Ref (MatchBind keyVar)) (TermGenerator pat)
      resultGen =
        DerivedFactGenerator pidRef (Ref (MatchVar keyVar)) (Tuple [])
      resultStmt2 =
        FlatStatement pidTy (Ref (MatchBind fidVar)) resultGen

    return
      ( FlatQuery result Nothing
          (mkStatementGroup (Floating resultStmt1 : Floating resultStmt2 : ord))
      , Nothing
      , retTy )

captureKey _ _ (FlatQuery _ Just{} _) _ =
  throwError "queries returning both a key and value are not supported"