{-
  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.
-}

-- | Collecting the variables mentioned by a term
module Glean.Query.Vars (
    VarId,
    VarSet,
    WhichVars(..),
    VarsOf(..),
    vars,
    varsBound,
    varsUsed,
    Fresh(..),
    fresh,
    freshWild,
    freshWildVars,
    reWild,
    reWildGenerator,
    reWildQuery,
  ) where

import qualified Data.IntSet as IntSet
import Data.IntSet (IntSet)
import qualified Data.IntMap as IntMap
import Data.IntMap (IntMap)
import Data.List.NonEmpty (NonEmpty)

import Glean.Display
import Glean.Query.Codegen.Types
import Glean.RTS.Term
import Glean.RTS.Types as RTS

type VarId = Int
type VarSet = IntSet
type VarMap = IntMap VarId

vars :: VarsOf a => a -> VarSet
vars x = varsOf AllVars x IntSet.empty

varsBound :: VarsOf a => a -> VarSet
varsBound x = varsOf VarsBound x IntSet.empty

varsUsed :: VarsOf a => a -> VarSet
varsUsed x = varsOf VarsUsed x IntSet.empty

data WhichVars = VarsUsed | VarsBound | AllVars
  deriving Eq

class VarsOf a where
  varsOf :: WhichVars -> a -> VarSet -> VarSet

instance VarsOf Generator where
  varsOf w (FactGenerator _ key val _) r = varsOf w key $! varsOf w val r
  varsOf w (TermGenerator exp) r = varsOf w exp r
  varsOf w (DerivedFactGenerator _ key val) r = varsOf w key $! varsOf w val r
  varsOf w (ArrayElementGenerator _ arr) r = varsOf w arr r
  varsOf w (SetElementGenerator _ set) r = varsOf w set r
  varsOf w (PrimCall _ args _) r = varsOf w args r

instance (VarsOf a) => VarsOf [a] where
  varsOf w container r = foldr (varsOf w) r container

instance (VarsOf a) => VarsOf (NonEmpty a) where
  varsOf w container r = foldr (varsOf w) r container

instance VarsOf m => VarsOf (Term m) where
  varsOf w t r = case t of
    Byte{} -> r
    Nat{} -> r
    ByteArray{} -> r
    String{} -> r
    Ref x -> varsOf w x r
    Tuple xs -> varsOf w xs r
    Array xs -> varsOf w xs r
    Alt _ x -> varsOf w x r

instance VarsOf (Match () Var) where
  varsOf w m r = case m of
    MatchWild{} -> r
    MatchNever{} -> r
    MatchFid{} -> r
    MatchBind (Var _ v _) -> if w == VarsUsed then r else IntSet.insert v r
    MatchVar (Var _ v _) -> if w == VarsBound then r else IntSet.insert v r
    MatchAnd a b -> varsOf w a $! varsOf w b r
    MatchPrefix _ t -> varsOf w t r
    MatchArrayPrefix _ty pre all -> varsOf w pre (varsOf w all r)
    MatchExt{} -> r

instance VarsOf CgQuery where
  varsOf w (CgQuery head stmts) r = varsOf w head $! varsOf w stmts r

instance VarsOf CgStatement where
  varsOf w (CgStatement lhs gen) r = varsOf w lhs $! varsOf w gen r
  varsOf w (CgAllStatement (Var _ v _) expr stmts) r =
    varsOf w expr $! varsOf w stmts $!
    if w /= VarsUsed then IntSet.insert v r else r
  varsOf w (CgNegation stmts) r = varsOf w stmts r
  varsOf w (CgDisjunction stmtss) r = varsOf w stmtss r
  varsOf w (CgConditional cond then_ else_) r =
    varsOf w cond $! varsOf w then_ $! varsOf w else_ r

-- -----------------------------------------------------------------------------
-- Fresh variables

class Fresh m where
  peek :: m VarId
  alloc :: m VarId

fresh :: (Monad m, Fresh m) => Type -> m Var
fresh ty = do
  n <- alloc
  return (Var ty n Nothing)

-- -----------------------------------------------------------------------------
-- Replace wildcards with fresh variables

freshWildVars :: (Monad m, Fresh m) => (a -> m a) -> a -> m (a, VarSet)
freshWildVars freshen thing = do
  v0 <- peek
  pat' <- freshen thing
  v1 <- peek
  return (pat', IntSet.fromList [v0..v1])

-- | Instantiate all the wildcards in a pattern with fresh
-- variables. This makes the pattern usable when we substitute it, for
-- two reasons: (1) it might occur in multiple places, and we must
-- ensure that it matches the same term in all places, and (2) it
-- might occur in an expression (E where ...) where wildcards don't
-- make sense.
freshWild :: forall m . (Monad m, Fresh m) => Pat -> m Pat
freshWild pat = mapM freshWildMatch pat
  where
  freshWildMatch :: Match a Var -> m (Match a Var)
  freshWildMatch m = case m of
    MatchWild ty -> MatchBind <$> fresh ty
    MatchPrefix str rest -> MatchPrefix str <$> mapM freshWildMatch rest
    MatchArrayPrefix ty pre all ->
      MatchArrayPrefix ty
        <$> (mapM.mapM) freshWildMatch pre
        <*> mapM freshWildMatch all
    MatchNever ty -> return (MatchNever ty)
    MatchFid f -> return (MatchFid f)
    MatchBind v -> return (MatchVar v)  -- also make all variables MatchVar
    MatchVar v -> return (MatchVar v)
    MatchAnd x y -> MatchAnd <$> mapM freshWildMatch x <*> mapM freshWildMatch y
    MatchExt _ -> error "freshWildMatch"

-- | Replace unused variables with wildcards and remap variable numbers
reWild :: VarMap -> Pat -> Pat
reWild used pat = fmap reWildMatch pat
  where
  reWildMatch :: Match () Var -> Match () Var
  reWildMatch m = case m of
    MatchWild ty -> MatchWild ty
    MatchBind (Var ty n x) -> case IntMap.lookup n used of
      Nothing -> MatchWild ty
      Just new -> MatchBind (Var ty new x)
    MatchVar (Var ty n x) -> case IntMap.lookup n used of
      Nothing -> error $ "reWild: " <> show (displayVerbose m)
      Just new -> MatchVar (Var ty new x)
    MatchPrefix str rest -> MatchPrefix str (fmap reWildMatch rest)
    MatchArrayPrefix ty pre all ->
      MatchArrayPrefix ty ((fmap . fmap) reWildMatch pre) (fmap reWildMatch all)
    MatchAnd x y -> MatchAnd (fmap reWildMatch x) (fmap reWildMatch y)
    MatchNever ty -> MatchNever ty
    MatchFid f -> MatchFid f
    MatchExt _ -> error "reWild"

reWildGenerator :: VarMap -> Generator -> Generator
reWildGenerator used gen = case gen of
  FactGenerator pid key val sec ->
    FactGenerator pid (reWild used key) (reWild used val) sec
  TermGenerator t -> TermGenerator (reWild used t)
  DerivedFactGenerator pid key val ->
    DerivedFactGenerator pid (reWild used key) (reWild used val)
  ArrayElementGenerator ty exp ->
    ArrayElementGenerator ty (reWild used exp)
  SetElementGenerator ty exp ->
    SetElementGenerator ty (reWild used exp)
  PrimCall op args ty ->
    PrimCall op (map (reWild used) args) ty

reWildStatement :: VarMap -> CgStatement -> CgStatement
reWildStatement used (CgStatement lhs rhs) =
  CgStatement (reWild used lhs) (reWildGenerator used rhs)
reWildStatement used s@(CgAllStatement (Var ty n x) expr stmts) =
  case IntMap.lookup n used of
    Nothing -> error $
      "reWildStatement: var " <> show n <>
      " not in scope in " <> show (displayVerbose s) <>
      "\nVarMap: " <> show used
    Just new ->
      CgAllStatement
        (Var ty new x) (reWild used expr)
        (map (reWildStatement used) stmts)
reWildStatement used (CgNegation stmts) =
  CgNegation (map (reWildStatement used) stmts)
reWildStatement used (CgDisjunction stmtss) =
  CgDisjunction (map (map (reWildStatement used)) stmtss)
reWildStatement used (CgConditional cond then_ else_) =
  CgConditional
    (map (reWildStatement used) cond)
    (map (reWildStatement used) then_)
    (map (reWildStatement used) else_)

reWildQuery :: VarMap -> CgQuery -> CgQuery
reWildQuery used (CgQuery head stmts) =
  CgQuery (reWild used head) (map (reWildStatement used) stmts)