    -- | Precompute free variables in a term (and store in 'ArgInfo').
    module Agda.TypeChecking.Free.Precompute
      ( PrecomputeFreeVars, precomputeFreeVars
      , precomputedFreeVars, precomputeFreeVars_ ) where

    import Control.Monad.Writer
    import Data.IntSet (IntSet)
    import qualified Data.IntSet as IntSet

    import Agda.Syntax.Common
    import Agda.Syntax.Internal

    type FV = Writer IntSet

    precomputeFreeVars_ :: PrecomputeFreeVars a => a -> a
    precomputeFreeVars_ :: forall a. PrecomputeFreeVars a => a -> a
    precomputeFreeVars_ = (a, IntSet) -> a
    forall a b. (a, b) -> a
    fst ((a, IntSet) -> a) -> (a -> (a, IntSet)) -> a -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Writer IntSet a -> (a, IntSet)
    forall w a. Writer w a -> (a, w)
    runWriter (Writer IntSet a -> (a, IntSet))
    -> (a -> Writer IntSet a) -> a -> (a, IntSet)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> Writer IntSet a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars

    precomputedFreeVars :: PrecomputeFreeVars a => a -> IntSet
    precomputedFreeVars :: forall a. PrecomputeFreeVars a => a -> IntSet
    precomputedFreeVars = (a, IntSet) -> IntSet
    forall a b. (a, b) -> b
    snd ((a, IntSet) -> IntSet) -> (a -> (a, IntSet)) -> a -> IntSet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Writer IntSet a -> (a, IntSet)
    forall w a. Writer w a -> (a, w)
    runWriter (Writer IntSet a -> (a, IntSet))
    -> (a -> Writer IntSet a) -> a -> (a, IntSet)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> Writer IntSet a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars

    class PrecomputeFreeVars a where
      precomputeFreeVars :: a -> FV a

      default precomputeFreeVars :: (Traversable c, PrecomputeFreeVars x, a ~ c x) => a -> FV a
      precomputeFreeVars = (x -> WriterT IntSet Identity x)
    -> c x -> WriterT IntSet Identity (c x)
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    traverse x -> WriterT IntSet Identity x
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars

    -- The instances where things actually happen: Arg, Abs and Term.

    maybePrecomputed :: PrecomputeFreeVars a => ArgInfo -> a -> FV (ArgInfo, a)
    maybePrecomputed :: forall a. PrecomputeFreeVars a => ArgInfo -> a -> FV (ArgInfo, a)
    maybePrecomputed ArgInfo
    i a
    x =
      case ArgInfo -> FreeVariables
    forall a. LensFreeVariables a => a -> FreeVariables
    getFreeVariables ArgInfo
    i of
        KnownFVs IntSet
    fv -> (ArgInfo
    i, a
    x) (ArgInfo, a) -> WriterT IntSet Identity () -> FV (ArgInfo, a)
    forall (f :: * -> *) a b. Functor f => a -> f b -> f a
    <$ IntSet -> WriterT IntSet Identity ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell IntSet
    fv
        FreeVariables
    UnknownFVs -> do
          (a
    x', IntSet
    fv) <- WriterT IntSet Identity a -> WriterT IntSet Identity (a, IntSet)
    forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
    listen (WriterT IntSet Identity a -> WriterT IntSet Identity (a, IntSet))
    -> WriterT IntSet Identity a -> WriterT IntSet Identity (a, IntSet)
    forall a b. (a -> b) -> a -> b
    $ a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    x
          (ArgInfo, a) -> FV (ArgInfo, a)
    forall (m :: * -> *) a. Monad m => a -> m a
    return (FreeVariables -> ArgInfo -> ArgInfo
    forall a. LensFreeVariables a => FreeVariables -> a -> a
    setFreeVariables (IntSet -> FreeVariables
    KnownFVs IntSet
    fv) ArgInfo
    i, a
    x')

    instance PrecomputeFreeVars a => PrecomputeFreeVars (Arg a) where
      precomputeFreeVars :: Arg a -> FV (Arg a)
    precomputeFreeVars arg :: Arg a
    arg@(Arg ArgInfo
    i a
    x) = (ArgInfo -> a -> Arg a) -> (ArgInfo, a) -> Arg a
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ArgInfo -> a -> Arg a
    forall e. ArgInfo -> e -> Arg e
    Arg ((ArgInfo, a) -> Arg a)
    -> WriterT IntSet Identity (ArgInfo, a) -> FV (Arg a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ArgInfo -> a -> WriterT IntSet Identity (ArgInfo, a)
    forall a. PrecomputeFreeVars a => ArgInfo -> a -> FV (ArgInfo, a)
    maybePrecomputed ArgInfo
    i a
    x

    -- Note that we don't store free variables in the Dom. The reason is that the
    -- ArgInfo in the Dom tends to get reused during type checking for the argument
    -- of that domain type, and it would be tedious and error prone to ensure that
    -- we don't accidentally inherit also the free variables. Moreover we don't
    -- really need the free variables of the Dom.
    instance PrecomputeFreeVars a => PrecomputeFreeVars (Dom a) where

    instance PrecomputeFreeVars a => PrecomputeFreeVars (Abs a) where
      precomputeFreeVars :: Abs a -> FV (Abs a)
    precomputeFreeVars (NoAbs ArgName
    x a
    b) = ArgName -> a -> Abs a
    forall a. ArgName -> a -> Abs a
    NoAbs ArgName
    x (a -> Abs a) -> WriterT IntSet Identity a -> FV (Abs a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    b
      precomputeFreeVars (Abs ArgName
    x a
    b) =
        (IntSet -> IntSet) -> FV (Abs a) -> FV (Abs a)
    forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
    censor ((Key -> Key) -> IntSet -> IntSet
    IntSet.map (Key -> Key -> Key
    forall a. Num a => a -> a -> a
    subtract Key
    1) (IntSet -> IntSet) -> (IntSet -> IntSet) -> IntSet -> IntSet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Key -> IntSet -> IntSet
    IntSet.delete Key
    0) (FV (Abs a) -> FV (Abs a)) -> FV (Abs a) -> FV (Abs a)
    forall a b. (a -> b) -> a -> b
    $
          ArgName -> a -> Abs a
    forall a. ArgName -> a -> Abs a
    Abs ArgName
    x (a -> Abs a) -> WriterT IntSet Identity a -> FV (Abs a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    b

    instance PrecomputeFreeVars Term where
      precomputeFreeVars :: Term -> FV Term
    precomputeFreeVars Term
    t =
        case Term
    t of
          Var Key
    x Elims
    es -> do
            IntSet -> WriterT IntSet Identity ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell (Key -> IntSet
    IntSet.singleton Key
    x)
            Key -> Elims -> Term
    Var Key
    x (Elims -> Term) -> WriterT IntSet Identity Elims -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          Lam ArgInfo
    i Abs Term
    b    -> ArgInfo -> Abs Term -> Term
    Lam ArgInfo
    i      (Abs Term -> Term) -> WriterT IntSet Identity (Abs Term) -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Abs Term -> WriterT IntSet Identity (Abs Term)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Abs Term
    b
          Lit{}      -> Term -> FV Term
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Term
    t
          Def QName
    f Elims
    es   -> QName -> Elims -> Term
    Def QName
    f      (Elims -> Term) -> WriterT IntSet Identity Elims -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          Con ConHead
    c ConInfo
    i Elims
    es -> ConHead -> ConInfo -> Elims -> Term
    Con ConHead
    c ConInfo
    i    (Elims -> Term) -> WriterT IntSet Identity Elims -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          Pi Dom Type
    a Abs Type
    b     -> (Dom Type -> Abs Type -> Term) -> (Dom Type, Abs Type) -> Term
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Dom Type -> Abs Type -> Term
    Pi ((Dom Type, Abs Type) -> Term)
    -> WriterT IntSet Identity (Dom Type, Abs Type) -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Dom Type, Abs Type)
    -> WriterT IntSet Identity (Dom Type, Abs Type)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars (Dom Type
    a, Abs Type
    b)
          Sort Sort
    s     -> Sort -> Term
    Sort       (Sort -> Term) -> WriterT IntSet Identity Sort -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Sort -> WriterT IntSet Identity Sort
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Sort
    s
          Level Level
    l    -> Level -> Term
    Level      (Level -> Term) -> WriterT IntSet Identity Level -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Level -> WriterT IntSet Identity Level
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Level
    l
          MetaV MetaId
    x Elims
    es -> MetaId -> Elims -> Term
    MetaV MetaId
    x    (Elims -> Term) -> WriterT IntSet Identity Elims -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          DontCare Term
    t -> Term -> Term
    DontCare   (Term -> Term) -> FV Term -> FV Term
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Term -> FV Term
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Term
    t
          Dummy{}    -> Term -> FV Term
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Term
    t

    -- The other instances are boilerplate.

    instance PrecomputeFreeVars Sort where
      precomputeFreeVars :: Sort -> WriterT IntSet Identity Sort
    precomputeFreeVars Sort
    s =
        case Sort
    s of
          Type Level
    a     -> Level -> Sort
    forall t. Level' t -> Sort' t
    Type (Level -> Sort)
    -> WriterT IntSet Identity Level -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Level -> WriterT IntSet Identity Level
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Level
    a
          Prop Level
    a     -> Level -> Sort
    forall t. Level' t -> Sort' t
    Prop (Level -> Sort)
    -> WriterT IntSet Identity Level -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Level -> WriterT IntSet Identity Level
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Level
    a
          Inf IsFibrant
    _ Integer
    _    -> Sort -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Sort
    s
          SSet Level
    a     -> Level -> Sort
    forall t. Level' t -> Sort' t
    SSet (Level -> Sort)
    -> WriterT IntSet Identity Level -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Level -> WriterT IntSet Identity Level
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Level
    a
          Sort
    SizeUniv   -> Sort -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Sort
    s
          Sort
    LockUniv   -> Sort -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Sort
    s
          PiSort Dom' Term Term
    a Sort
    s1 Abs Sort
    s2 -> Dom' Term Term -> Sort -> Abs Sort -> Sort
    forall t. Dom' t t -> Sort' t -> Abs (Sort' t) -> Sort' t
    PiSort (Dom' Term Term -> Sort -> Abs Sort -> Sort)
    -> WriterT IntSet Identity (Dom' Term Term)
    -> WriterT IntSet Identity (Sort -> Abs Sort -> Sort)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Dom' Term Term -> WriterT IntSet Identity (Dom' Term Term)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Dom' Term Term
    a WriterT IntSet Identity (Sort -> Abs Sort -> Sort)
    -> WriterT IntSet Identity Sort
    -> WriterT IntSet Identity (Abs Sort -> Sort)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Sort -> WriterT IntSet Identity Sort
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Sort
    s1 WriterT IntSet Identity (Abs Sort -> Sort)
    -> WriterT IntSet Identity (Abs Sort)
    -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Abs Sort -> WriterT IntSet Identity (Abs Sort)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Abs Sort
    s2
          FunSort Sort
    s1 Sort
    s2 -> (Sort -> Sort -> Sort) -> (Sort, Sort) -> Sort
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Sort -> Sort -> Sort
    forall t. Sort' t -> Sort' t -> Sort' t
    FunSort ((Sort, Sort) -> Sort)
    -> WriterT IntSet Identity (Sort, Sort)
    -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Sort, Sort) -> WriterT IntSet Identity (Sort, Sort)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars (Sort
    s1, Sort
    s2)
          UnivSort Sort
    s -> Sort -> Sort
    forall t. Sort' t -> Sort' t
    UnivSort (Sort -> Sort)
    -> WriterT IntSet Identity Sort -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Sort -> WriterT IntSet Identity Sort
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Sort
    s
          MetaS MetaId
    x Elims
    es -> MetaId -> Elims -> Sort
    forall t. MetaId -> [Elim' t] -> Sort' t
    MetaS MetaId
    x (Elims -> Sort)
    -> WriterT IntSet Identity Elims -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          DefS QName
    d Elims
    es  -> QName -> Elims -> Sort
    forall t. QName -> [Elim' t] -> Sort' t
    DefS QName
    d (Elims -> Sort)
    -> WriterT IntSet Identity Elims -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Elims -> WriterT IntSet Identity Elims
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Elims
    es
          DummyS{}   -> Sort -> WriterT IntSet Identity Sort
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Sort
    s

    instance PrecomputeFreeVars Level where
      precomputeFreeVars :: Level -> WriterT IntSet Identity Level
    precomputeFreeVars (Max Integer
    n [PlusLevel' Term]
    ls) = Integer -> [PlusLevel' Term] -> Level
    forall t. Integer -> [PlusLevel' t] -> Level' t
    Max Integer
    n ([PlusLevel' Term] -> Level)
    -> WriterT IntSet Identity [PlusLevel' Term]
    -> WriterT IntSet Identity Level
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [PlusLevel' Term] -> WriterT IntSet Identity [PlusLevel' Term]
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars [PlusLevel' Term]
    ls

    instance PrecomputeFreeVars PlusLevel where
      precomputeFreeVars :: PlusLevel' Term -> FV (PlusLevel' Term)
    precomputeFreeVars (Plus Integer
    n Term
    l) = Integer -> Term -> PlusLevel' Term
    forall t. Integer -> t -> PlusLevel' t
    Plus Integer
    n (Term -> PlusLevel' Term) -> FV Term -> FV (PlusLevel' Term)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Term -> FV Term
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Term
    l

    instance PrecomputeFreeVars Type where
      precomputeFreeVars :: Type -> FV Type
    precomputeFreeVars (El Sort
    s Term
    t) = (Sort -> Term -> Type) -> (Sort, Term) -> Type
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Sort -> Term -> Type
    forall t a. Sort' t -> a -> Type'' t a
    El ((Sort, Term) -> Type)
    -> WriterT IntSet Identity (Sort, Term) -> FV Type
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Sort, Term) -> WriterT IntSet Identity (Sort, Term)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars (Sort
    s, Term
    t)

    -- Note: don't use default instance, since that bypasses the 'Arg' in 'Apply'.
    instance PrecomputeFreeVars a => PrecomputeFreeVars (Elim' a) where
      precomputeFreeVars :: Elim' a -> FV (Elim' a)
    precomputeFreeVars Elim' a
    e =
        case Elim' a
    e of
          Apply Arg a
    x      -> Arg a -> Elim' a
    forall a. Arg a -> Elim' a
    Apply (Arg a -> Elim' a)
    -> WriterT IntSet Identity (Arg a) -> FV (Elim' a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Arg a -> WriterT IntSet Identity (Arg a)
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars Arg a
    x
          IApply a
    a a
    x a
    y -> a -> a -> a -> Elim' a
    forall a. a -> a -> a -> Elim' a
    IApply (a -> a -> a -> Elim' a)
    -> WriterT IntSet Identity a
    -> WriterT IntSet Identity (a -> a -> Elim' a)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    a WriterT IntSet Identity (a -> a -> Elim' a)
    -> WriterT IntSet Identity a
    -> WriterT IntSet Identity (a -> Elim' a)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    x WriterT IntSet Identity (a -> Elim' a)
    -> WriterT IntSet Identity a -> FV (Elim' a)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    y
          Proj{}       -> Elim' a -> FV (Elim' a)
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Elim' a
    e

    -- The very boilerplate instances

    instance PrecomputeFreeVars a => PrecomputeFreeVars [a] where
    instance PrecomputeFreeVars a => PrecomputeFreeVars (Maybe a) where

    instance (PrecomputeFreeVars a, PrecomputeFreeVars b) => PrecomputeFreeVars (a, b) where
      precomputeFreeVars :: (a, b) -> FV (a, b)
    precomputeFreeVars (a
    x, b
    y) = (,) (a -> b -> (a, b))
    -> WriterT IntSet Identity a
    -> WriterT IntSet Identity (b -> (a, b))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> a -> WriterT IntSet Identity a
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars a
    x WriterT IntSet Identity (b -> (a, b))
    -> WriterT IntSet Identity b -> FV (a, b)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> b -> WriterT IntSet Identity b
    forall a. PrecomputeFreeVars a => a -> FV a
    precomputeFreeVars b
    y