{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -Wunused-imports #-}

#if  __GLASGOW_HASKELL__ > 902
{-# OPTIONS_GHC -fworker-wrapper-cbv #-}
#endif

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

import Agda.Utils.StrictWriter ( Writer, runWriter, censor, listen, tell )

import Agda.Utils.Singleton
import Agda.Utils.VarSet (VarSet)
import qualified Agda.Utils.VarSet as VarSet

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

type FV = Writer VarSet

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

precomputedFreeVars :: PrecomputeFreeVars a => a -> VarSet
precomputedFreeVars :: forall a. PrecomputeFreeVars a => a -> VarSet
precomputedFreeVars = (a, VarSet) -> VarSet
forall a b. (a, b) -> b
snd ((a, VarSet) -> VarSet) -> (a -> (a, VarSet)) -> a -> VarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer VarSet a -> (a, VarSet)
forall w a. Monoid w => Writer w a -> (a, w)
runWriter (Writer VarSet a -> (a, VarSet))
-> (a -> Writer VarSet a) -> a -> (a, VarSet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Writer VarSet 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 -> Writer VarSet x) -> c x -> Writer VarSet (c x)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> c a -> f (c b)
traverse x -> Writer VarSet 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 VarSet
fv -> (ArgInfo
i, a
x) (ArgInfo, a) -> Writer VarSet () -> FV (ArgInfo, a)
forall a b. a -> Writer VarSet b -> Writer VarSet a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ VarSet -> Writer VarSet ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell VarSet
fv
    FreeVariables
UnknownFVs -> do
      (x', fv) <- Writer VarSet a -> Writer VarSet (a, VarSet)
forall a. Writer VarSet a -> Writer VarSet (a, VarSet)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (Writer VarSet a -> Writer VarSet (a, VarSet))
-> Writer VarSet a -> Writer VarSet (a, VarSet)
forall a b. (a -> b) -> a -> b
$ a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x
      return (setFreeVariables (KnownFVs fv) i, 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) -> Writer VarSet (ArgInfo, a) -> FV (Arg a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ArgInfo -> a -> Writer VarSet (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) -> Writer VarSet a -> FV (Abs a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
b
  precomputeFreeVars (Abs ArgName
x a
b) =
    (VarSet -> VarSet) -> FV (Abs a) -> FV (Abs a)
forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
censor (Int -> VarSet -> VarSet
VarSet.strengthen Int
1) (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) -> Writer VarSet a -> FV (Abs a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Writer VarSet 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 Int
x [Elim]
es -> do
        VarSet -> Writer VarSet ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Int -> VarSet
forall el coll. Singleton el coll => el -> coll
singleton Int
x)
        Int -> [Elim] -> Term
Var Int
x ([Elim] -> Term) -> Writer VarSet [Elim] -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Elim] -> Writer VarSet [Elim]
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars [Elim]
es
      Lam ArgInfo
i Abs Term
b    -> ArgInfo -> Abs Term -> Term
Lam ArgInfo
i      (Abs Term -> Term) -> Writer VarSet (Abs Term) -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Abs Term -> Writer VarSet (Abs Term)
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Abs Term
b
      Lit{}      -> Term -> FV Term
forall a. a -> Writer VarSet a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term
t
      Def QName
f [Elim]
es   -> QName -> [Elim] -> Term
Def QName
f      ([Elim] -> Term) -> Writer VarSet [Elim] -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Elim] -> Writer VarSet [Elim]
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars [Elim]
es
      Con ConHead
c ConInfo
i [Elim]
es -> ConHead -> ConInfo -> [Elim] -> Term
Con ConHead
c ConInfo
i    ([Elim] -> Term) -> Writer VarSet [Elim] -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Elim] -> Writer VarSet [Elim]
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars [Elim]
es
      Pi Dom' Term Type
a Abs Type
b     -> (Dom' Term Type -> Abs Type -> Term)
-> (Dom' Term Type, Abs Type) -> Term
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Dom' Term Type -> Abs Type -> Term
Pi ((Dom' Term Type, Abs Type) -> Term)
-> Writer VarSet (Dom' Term Type, Abs Type) -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Dom' Term Type, Abs Type)
-> Writer VarSet (Dom' Term Type, Abs Type)
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars (Dom' Term Type
a, Abs Type
b)
      Sort Sort' Term
s     -> Sort' Term -> Term
Sort       (Sort' Term -> Term) -> Writer VarSet (Sort' Term) -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sort' Term -> Writer VarSet (Sort' Term)
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Sort' Term
s
      Level Level' Term
l    -> Level' Term -> Term
Level      (Level' Term -> Term) -> Writer VarSet (Level' Term) -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Level' Term -> Writer VarSet (Level' Term)
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars Level' Term
l
      MetaV MetaId
x [Elim]
es -> MetaId -> [Elim] -> Term
MetaV MetaId
x    ([Elim] -> Term) -> Writer VarSet [Elim] -> FV Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Elim] -> Writer VarSet [Elim]
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars [Elim]
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 a. a -> Writer VarSet a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term
t

-- The other instances are boilerplate.

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

instance PrecomputeFreeVars Level where
  precomputeFreeVars :: Level' Term -> Writer VarSet (Level' Term)
precomputeFreeVars (Max Integer
n [PlusLevel' Term]
ls) = Integer -> [PlusLevel' Term] -> Level' Term
forall t. Integer -> [PlusLevel' t] -> Level' t
Max Integer
n ([PlusLevel' Term] -> Level' Term)
-> Writer VarSet [PlusLevel' Term] -> Writer VarSet (Level' Term)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [PlusLevel' Term] -> Writer VarSet [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' Term
s Term
t) = (Sort' Term -> Term -> Type) -> (Sort' Term, Term) -> Type
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Sort' Term -> Term -> Type
forall t a. Sort' t -> a -> Type'' t a
El ((Sort' Term, Term) -> Type)
-> Writer VarSet (Sort' Term, Term) -> FV Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Sort' Term, Term) -> Writer VarSet (Sort' Term, Term)
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars (Sort' Term
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) -> Writer VarSet (Arg a) -> FV (Elim' a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Arg a -> Writer VarSet (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)
-> Writer VarSet a -> Writer VarSet (a -> a -> Elim' a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
a Writer VarSet (a -> a -> Elim' a)
-> Writer VarSet a -> Writer VarSet (a -> Elim' a)
forall a b.
Writer VarSet (a -> b) -> Writer VarSet a -> Writer VarSet b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x Writer VarSet (a -> Elim' a) -> Writer VarSet a -> FV (Elim' a)
forall a b.
Writer VarSet (a -> b) -> Writer VarSet a -> Writer VarSet b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
y
      Proj{}       -> Elim' a -> FV (Elim' a)
forall a. a -> Writer VarSet 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))
-> Writer VarSet a -> Writer VarSet (b -> (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Writer VarSet a
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars a
x Writer VarSet (b -> (a, b)) -> Writer VarSet b -> FV (a, b)
forall a b.
Writer VarSet (a -> b) -> Writer VarSet a -> Writer VarSet b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> b -> Writer VarSet b
forall a. PrecomputeFreeVars a => a -> FV a
precomputeFreeVars b
y