-- | 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_ = forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall w a. Writer w a -> (a, w)
runWriter forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars

precomputedFreeVars :: PrecomputeFreeVars a => a -> IntSet
precomputedFreeVars :: forall a. PrecomputeFreeVars a => a -> IntSet
precomputedFreeVars = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall w a. Writer w a -> (a, w)
runWriter forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse 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 forall a. LensFreeVariables a => a -> FreeVariables
getFreeVariables ArgInfo
i of
    KnownFVs IntSet
fv -> (ArgInfo
i, a
x) forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell IntSet
fv
    FreeVariables
UnknownFVs -> do
      (a
x', IntSet
fv) <- forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen forall a b. (a -> b) -> a -> b
$ forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x
      forall (m :: * -> *) a. Monad m => a -> m a
return (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) = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall e. ArgInfo -> e -> Arg e
Arg forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) = forall a. ArgName -> a -> Abs a
NoAbs ArgName
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
b
  precomputeFreeVars (Abs ArgName
x a
b) =
    forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
censor ((Key -> Key) -> IntSet -> IntSet
IntSet.map (forall a. Num a => a -> a -> a
subtract Key
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> IntSet -> IntSet
IntSet.delete Key
0) forall a b. (a -> b) -> a -> b
$
      forall a. ArgName -> a -> Abs a
Abs ArgName
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
        forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Key -> IntSet
IntSet.singleton Key
x)
        Key -> Elims -> Term
Var Key
x forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Elims
es
      Lam ArgInfo
i Abs Term
b    -> ArgInfo -> Abs Term -> Term
Lam ArgInfo
i      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Abs Term
b
      Lit{}      -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Term
t
      Def QName
f Elims
es   -> QName -> Elims -> Term
Def QName
f      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Elims
es
      Pi Dom Type
a Abs Type
b     -> forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Dom Type -> Abs Type -> Term
Pi forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars (Dom Type
a, Abs Type
b)
      Sort Sort
s     -> Sort -> Term
Sort       forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Sort
s
      Level Level
l    -> Level -> Term
Level      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Level
l
      MetaV MetaId
x Elims
es -> MetaId -> Elims -> Term
MetaV MetaId
x    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Elims
es
      DontCare Term
t -> Term -> Term
DontCare   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Term
t
      Dummy{}    -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Term
t

-- The other instances are boilerplate.

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

instance PrecomputeFreeVars Level where
  precomputeFreeVars :: Level -> FV Level
precomputeFreeVars (Max Integer
n [PlusLevel' Term]
ls) = forall t. Integer -> [PlusLevel' t] -> Level' t
Max Integer
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) = forall t. Integer -> t -> PlusLevel' t
Plus Integer
n forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Term
l

instance PrecomputeFreeVars Type where
  precomputeFreeVars :: Type -> FV Type
precomputeFreeVars (El Sort
s Term
t) = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall t a. Sort' t -> a -> Type'' t a
El forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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      -> forall a. Arg a -> Elim' a
Apply forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Arg a
x
      IApply a
a a
x a
y -> forall a. a -> a -> a -> Elim' a
IApply forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
y
      Proj{}       -> 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) = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars b
y