content: -- | Imperfect discrimination trees for indexing data by internal
  -- syntax.
  module Agda.TypeChecking.DiscrimTree
    ( insertDT
    , lookupDT, lookupUnifyDT, QueryResult(..)
    , deleteFromDT
    )
    where
  
  import qualified Data.Map.Strict as Map
  import qualified Data.Set as Set
  import Data.Foldable
  import Data.Maybe
  
  import Control.Monad.Trans.Maybe
  import Control.Monad.Trans
  import Control.Monad
  
  import Agda.Syntax.Internal
  import Agda.Syntax.Common
  
  import Agda.TypeChecking.Substitute
  import Agda.TypeChecking.Telescope
  import Agda.TypeChecking.Records
  import Agda.TypeChecking.Reduce
  import Agda.TypeChecking.Pretty
  import Agda.TypeChecking.Monad
  import Agda.TypeChecking.Free
  
  import Agda.TypeChecking.DiscrimTree.Types
  
  import qualified Agda.Interaction.Options.ProfileOptions as Profile
  
  import Agda.Utils.Impossible
  import Agda.Utils.Trie (Trie(..))
  
  -- | Dummy term to use as a stand-in for expanded eta-records while
  -- building instance trees.
  etaExpansionDummy :: Term
  etaExpansionDummy = Dummy "eta-record argument in instance head" []
  
  -- | Extract a list of arguments from the list of eliminations; If
  -- called while *adding* an instance, additionally replace any arguments
  -- that might belong to an eta-record by dummy terms.
  termKeyElims
    :: Bool      -- ^ Are we adding or looking up an instance?
    -> TCM Type  -- ^ Continuation to compute the type of the arguments in the spine.
    -> [Arg Term] -- ^ The spine.
    -> TCM (Int, [Term])
  
  -- Since the case tree was generated with wildcards everywhere an eta
  -- record appeared, if we're *looking up* an instance, we don't have to
  -- do the censorship again.
  termKeyElims False _ es = pure (length es, map unArg es)
  
  termKeyElims precise ty args = do
    let
      go ty (Arg _ a:as) = flip (ifPiTypeB ty) (patternViolation . getBlocker) \dom ty' -> do
  
        -- Is this argument an eta record type --- or a blocked value? In
        -- either case, we replace this position by a dummy, to make sure
        -- that eta-equality is respected.
        maybeEta <- ifBlocked (unDom dom) (\_ _ -> pure True) \_ tm ->
          isJust <$> isEtaRecordType tm
  
        let
          here
            | maybeEta  = etaExpansionDummy
            | otherwise = a
  
        (k, there) <- addContext dom (go (unAbs ty') as)
        pure (k + 1, here:there)
  
      go _ [] = pure (0, [])
  
    ty >>= flip go args
  
  -- | Ticky profiling for the reason behind "inexactness" in instance
  -- search. If at some point while narrowing the set of candidates we had
  -- to go through all the possibilities, one of these counters is
  -- incremented.
  tickExplore :: Term -> TCM ()
  tickExplore tm = whenProfile Profile.Instances do
    tick "explore"
  
    case tm of
      Def{}      -> tick "explore: Def"
      Var{}      -> tick "explore: Var"
      Lam _ v
        -- These two are a hunch: just like FunK, it might be worth
        -- optimising for the case where a lambda is constant (which is
        -- easy to handle, by just pretending the term is something else).
        -- These would come up in e.g. Dec (PathP (λ i → Nat) x y)
        | NoAbs{} <- v -> tick "explore: constant function"
        | Abs _ b <- v, not (0 `freeIn` b) -> tick "explore: constant function"
  
        | otherwise    -> tick "explore: Lam"
      Lit{}      -> tick "explore: Lit"
      Sort{}     -> tick "explore: Sort"
      Level{}    -> tick "explore: Level"
      MetaV{}    -> tick "explore: Meta"
      DontCare{} -> tick "explore: DontCare"
      _ -> pure ()
  
  -- | Split a term into a 'Key' and some arguments. The 'Key' indicates
  -- whether or not the 'Term' is in head-normal form, and provides a
  -- quick way to match on the head.
  --
  -- The 'Int' argument indicates how free a variable must be to be
  -- considered a 'LocalK'.
  --
  -- Presently, non-head-normal terms end up with an empty argument list.
  splitTermKey :: Bool -> Int -> Term -> TCM (Key, [Term], Blocker)
  splitTermKey precise local tm = catchPatternErr (\b -> pure (FlexK, [], b)) do
    (b, tm') <- ifBlocked tm (\b _ -> patternViolation b) (\b -> fmap (b,) . constructorForm)
  
    case tm' of
      -- Adding a 'Def' to the key poses a few problems when opacity (or
      -- abstractness) are involved (see issue #7304). Suppose we have an
      -- opaque binding `X = Y`, and an opaque instance `C X`. The problem
      -- is as follows:
      --
      --   if we unfold X → Y when adding the instance, then it will not
      --   get recorded as an instance for C X, only C Y; this is 7304b.
      --
      --   if we *don't* unfold X → Y, then it only gets added as an
      --   instance of C X; in opaque blocks where X is allowed to unfold,
      --   we *won't* find it, because we're looking for C Y.
      --
      -- The solution is to throw our hands up and say "not our problem".
      -- The discrimination tree is allowed to return more results than
      -- strictly necessary, after all, so the solution is to add an
      -- instance for *neither* of C X or C Y, but instead, to treat all
      -- 'Def's headed by 'AbstractDefn' as though they were flexible
      -- (think "as though they were metas").
      Def q as | ReallyNotBlocked <- b, (as, _) <- splitApplyElims as -> do
        info <- getConstInfo q
        case theDef info of
          AbstractDefn{} | precise -> pure (FlexK, [], neverUnblock)
          _ -> do
            (arity, as) <- termKeyElims precise (pure (defType info)) as
            pure (RigidK q arity, as, neverUnblock)
  
      -- When adding a quantified instance, we record how many 'Pi's we went
      -- under, and only variables beyond those are considered LocalK. The
      -- others are considered FlexK since they're "pattern variables" of
      -- the instance.
      Var i as | i >= local, Just as <- allApplyElims as -> do
        let ty = unDom <$> domOfBV i
        (arity, as) <- termKeyElims precise ty as
        pure (LocalK (i - local) arity, as, neverUnblock)
  
      -- When looking up an instance, it's better to treat variables and
      -- neutral definitions as rigid things regardless of their spines
      -- (especially if they have projections), than it is to try to
      -- represent them accurately.
      Def q as | not precise             -> pure (RigidK q 0, [], neverUnblock)
      Var i as | not precise, i >= local -> pure (LocalK (i - local) 0, [], neverUnblock)
  
      Con ch _ as | Just as <- allApplyElims as -> do
        let
          q  = conName ch
          ty = defType <$> getConstInfo q
        (arity, as) <- termKeyElims precise ty as
        pure (RigidK q arity, as, neverUnblock)
  
      Pi dom ret ->
        let
          -- If we're looking at a non-dependent function type, then we
          -- might as well represent the codomain accurately; Otherwise,
          -- turn the codomain into a wildcard.
          --
          -- The use of a dummy term *shouldn't* leak to the user, because
          -- when we call splitTermKey again, it'll be handled by the last
          -- case, and become a FlexK.
          ret' = case isNoAbs (unEl <$> ret) of
            Just b  -> b
            Nothing -> __DUMMY_TERM__
        in pure (PiK (getHiding dom), [unEl (unDom dom), ret'], neverUnblock)
  
      Lam _ body
        -- Constant lambdas come up quite a bit, particularly (in cubical
        -- mode) as the domain of a PathP. Having this trick improves the
        -- indexing of 'Dec' instances in the 1Lab significantly.
        | Just b <- isNoAbs body -> pure (ConstK, [b], neverUnblock)
  
      -- Probably not a good idea for accurate indexing if universes
      -- overlap literally everything else.
      Sort _ -> pure (SortK, [], neverUnblock)
  
      _ -> do
        reportSDoc "tc.instance.split" 30 $ pretty tm
        pure (FlexK, [], neverUnblock)
  
  termPath :: Int -> Bool -> Int -> [Key] -> [Term] -> TCM [Key]
  termPath 0 _ _ acc _                        = pure $! reverse acc
  termPath limit toplevel local acc []        = pure $! reverse acc
  termPath limit toplevel local acc (tm:todo) = do
  
    -- We still want to ignore abstractness at the very top-level of
    -- instance heads, for issue #6941, to ensure that each instance ends
    -- up in the right 'class'. See the comment in `splitTermKey` about
    -- abstract definitions.
    (k, as, blk) <-
      if toplevel
        then ignoreAbstractMode (splitTermKey True local tm)
        else splitTermKey True local tm
  
    reportSDoc "tc.instance.discrim.add" 666 $ vcat
      [ "k:  " <+> prettyTCM k
      , "as: " <+> prettyTCM as
      , "blk:" <+> prettyTCM blk
      , "lim:" <+> prettyTCM limit
      ]
    termPath (limit - 1) False local (k:acc) (as <> todo)
  
  -- | Maximum length for the keys (thus, depth for the discrimination
  -- tree) that should be used.
  --
  -- Adding an instance still causes reduction of everything 'rigid' in
  -- the type of the instance (so if normalisation is slow, so will be
  -- adding the instance) but this limit prevents us from building a
  -- discrimination tree with 2^32 intermediate @case 0 of suc → ...@
  -- nodes.
  discrimTreeDepthLimit :: Int
  discrimTreeDepthLimit = 16
  
  -- | Insert a value into the discrimination tree, turning variables into
  -- rigid locals or wildcards depending on the given scope.
  insertDT
    :: (Ord a, PrettyTCM a)
    => Int   -- ^ Number of variables to consider wildcards, e.g. the number of leading invisible pis in an instance type.
    -> Term  -- ^ The term to use as a key
    -> a
    -> DiscrimTree a
    -> TCM (DiscrimTree a)
  insertDT local key val tree = do
    path <- termPath discrimTreeDepthLimit True local [] [key]
    let it = singletonDT path val
    reportSDoc "tc.instance.discrim.add" 20 $ vcat
      [ "added value" <+> prettyTCM val <+> "to discrimination tree with case"
      , nest 2 (prettyTCM it)
      , "its type:"
      , nest 2 (prettyTCM key)
      , "its path:"
      , nest 2 (prettyTCM path)
      ]
    pure $ mergeDT it tree
  
  -- | If a term matches this key, how many arguments does it place on the
  -- spine?
  keyArity :: Key -> Int
  keyArity = \case
    RigidK _ a -> a
    LocalK _ a -> a
    PiK _      -> 2
    ConstK     -> 1
    SortK      -> 0
    FlexK      -> 0
  
  data QueryResult a = QueryResult
    { resultValues  :: Set.Set a
    , resultBlocker :: Blocker
    }
  
  instance Ord a => Semigroup (QueryResult a) where
    QueryResult s b <> QueryResult s' b' = QueryResult (s <> s') (b `unblockOnEither` b')
  
  instance Ord a => Monoid (QueryResult a) where
    mempty = QueryResult mempty neverUnblock
  
  setResult :: Set.Set a -> QueryResult a
  setResult = flip QueryResult neverUnblock
  
  blockerResult :: Blocker -> QueryResult a
  blockerResult = QueryResult Set.empty
  
  -- | Look up a 'Term' in the given discrimination tree, treating local
  -- variables as rigid symbols. The returned set is guaranteed to contain
  -- everything that could overlap the given key.
  lookupDT :: forall a. (Ord a, PrettyTCM a) => Term -> DiscrimTree a -> TCM (QueryResult a)
  lookupDT = lookupDT' True
  
  -- | Look up a 'Term' in the given discrimination tree, treating local
  -- variables as wildcards.
  lookupUnifyDT :: forall a. (Ord a, PrettyTCM a) => Term -> DiscrimTree a -> TCM (QueryResult a)
  lookupUnifyDT = lookupDT' False
  
  lookupDT'
    :: forall a. (Ord a, PrettyTCM a)
    => Bool -- ^ Should local variables be treated as rigid?
    -> Term -- ^ The term to use as key
    -> DiscrimTree a
    -> TCM (QueryResult a)
  lookupDT' localsRigid term tree = match True [term] tree where
  
    split :: Term -> TCM (Key, [Term], Blocker)
    split tm | localsRigid = splitTermKey False 0 tm
    split tm = do
      ctx <- getContextSize
      splitTermKey False ctx tm
  
    ignoreAbstractMaybe :: forall a. Bool -> TCM a -> TCM a
    ignoreAbstractMaybe True  = ignoreAbstractMode
    ignoreAbstractMaybe False = id
  
    -- Match a spine against *all* clauses.
    explore :: [Term] -> [Term] -> [Term] -> [(Key, DiscrimTree a)] -> TCM (QueryResult a)
    explore sp0 sp1 args bs = do
      let
        cont (key, trie) res = do
          -- At the moment, explore will always be called with empty args.
          -- But even if this restriction is lifted in the future, we have
          -- to be careful about exploring. Consider:
          --
          --   instance
          --     _ : Foo (con x)
          --
          --   ⊢ Foo ?0
          --
          -- Since ?0 might be applied to more or less arguments than the
          -- one argument that is expected to be between sp0 and sp1 after
          -- matching con, we need to make sure that the spine has the
          -- right number of arguments, otherwise the (sp0, t:sp1) pattern
          -- for a Case will fail.
          let
            dummy n = Dummy ("_pad" <> show n) []
            args' = take (keyArity key) (args ++ [ dummy n | n <- [0..] ])
  
          reportSDoc "tc.instance.discrim.lookup" 99 $ vcat
            [ "explore" <+> prettyTCM key <+> pretty (keyArity key) <+> pretty (length args)
            , nest 2 (prettyTCM trie)
            , "sp0:  " <+> prettyTCM sp0
            , "sp1:  " <+> prettyTCM sp1
            , "args: " <+> prettyTCM args
            , "args':" <+> prettyTCM args'
            ]
          (<> res) <$> match False (sp0 ++ args' ++ sp1) trie
  
      foldrM cont mempty bs
  
    match :: Bool -> [Term] -> DiscrimTree a -> TCM (QueryResult a)
    match toplevel ts EmptyDT    = pure mempty
    match toplevel ts (DoneDT t) = setResult t <$ do
      reportSDoc "tc.instance.discrim.lookup" 99 $ vcat
        [ "done" <+> prettyTCM ts
        , "  →" <+> prettyTCM t
        ]
  
    match toplevel ts tree@(CaseDT i branches rest) | (sp0, t:sp1) <- splitAt i ts = do
      let
        (sp0, t:sp1) = splitAt i ts
        visit k sp' = case Map.lookup k branches of
          Just m  -> match False sp' m
          Nothing -> pure mempty
  
      unless toplevel $ reportSDoc "tc.instance.discrim.lookup" 99 $ vcat
        [ "match" <+> prettyTCM sp0 <+> ("«" <> prettyTCM t <> "»") <+> prettyTCM sp1
        , prettyTCM tree
        ]
  
      -- TODO (Amy, 2024-02-12): Could use reduceB in splitTermKey, and
      -- the blocker here, to suspend instances more precisely when there
      -- is an ambiguity.
      ignoreAbstractMaybe toplevel (split t) >>= \case
        (FlexK, args, blocker) -> do
  
          reportSDoc "tc.instance.discrim.lookup" 99 $ vcat
            [ "flexible term was forced"
            , "t:" <+> (pretty =<< instantiate t)
            , "will explore" <+> pretty (length branches + 1) <+> "branches"
            ]
          tickExplore t
  
          -- If we have a "flexible head" at this position then instance
          -- search *at this point* degenerates to looking for all
          -- possible matches.
          --
          -- In any nested CaseDTs, however, it's possible for us to
          -- recover and go back to productively matching. Consider:
          --
          --    instance
          --      xa : X T1 A
          --      xb : X T2 B
          --
          --    ⊢ X ?0 A
          --
          -- Since ?0 is way too flabby to narrow which of T1 or T2 should
          -- be taken, we take both. But then we match A against A and B:
          -- this query will only return {xa}.
  
          branches <- explore sp0 sp1 args $ Map.toList branches
          rest <- match False ts rest
  
          pure $! rest <> branches <> blockerResult blocker
  
        (k, args, blocker) -> do
          let sp' = sp0 ++ args ++ sp1
  
          -- Actually take the branch corresponding to our rigid head.
          branch <- visit k sp'
  
          -- When exploring the rest of the tree, the value we cased on
          -- has to be put back in the tree. mergeDT does not perform
          -- commuting conversions to ensure that variables aren't
          -- repeatedly cased on.
          rest <- match False ts rest
  
          pure $! rest <> branch
  
    match _ ts tree@(CaseDT i _ rest) = do
      reportSDoc "tc.instance.discrim.lookup" 99 $ vcat
        [ "IMPOSSIBLE match" <+> prettyTCM ts
        , prettyTCM tree
        ]
      -- This really is impossible: since each branch is annotated with
      -- its arity, we only take branches corresponding to neutrals which
      -- exploded into enough arguments.
      __IMPOSSIBLE__
  
  -- | Smart constructor for a leaf node.
  doneDT :: Set.Set a -> DiscrimTree a
  doneDT s | Set.null s = EmptyDT
  doneDT s = DoneDT s
  
  -- | Remove a set of values from the discrimination tree. The tree is
  -- rebuilt so that cases with no leaves are removed.
  deleteFromDT :: Ord a => Set.Set a -> DiscrimTree a -> DiscrimTree a
  deleteFromDT gone = \case
    EmptyDT -> EmptyDT
    DoneDT s -> doneDT $! Set.difference s gone
    CaseDT i s k ->
      let
        del x = case deleteFromDT gone x of
          EmptyDT -> Nothing
          dt'     -> Just dt'
  
        s' = Map.mapMaybe del s
        k' = deleteFromDT gone k
      in if | Map.null s' -> k'
            | otherwise   -> CaseDT i s' k'
  