-- | 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