    {-# LANGUAGE GeneralizedNewtypeDeriving #-}
    -- | Computing the free variables of a term.
    --
    -- The distinction between rigid and strongly rigid occurrences comes from:
    --   Jason C. Reed, PhD thesis, 2009, page 96 (see also his LFMTP 2009 paper)
    --
    -- The main idea is that x = t(x) is unsolvable if x occurs strongly rigidly
    -- in t.  It might have a solution if the occurrence is not strongly rigid, e.g.
    --
    --   x = \f -> suc (f (x (\ y -> k)))  has  x = \f -> suc (f (suc k))
    --
    -- [Jason C. Reed, PhD thesis, page 106]
    --
    -- Under coinductive constructors, occurrences are never strongly rigid.
    -- Also, function types and lambdas do not establish strong rigidity.
    -- Only inductive constructors do so.
    -- (See issue 1271).

    module Agda.TypeChecking.Free
        ( FreeVars(..)
        , VarCounts(..)
        , Free
        , IsVarSet(..)
        , IgnoreSorts(..)
        , runFree , rigidVars, relevantVars, allVars
        , allFreeVars, allFreeVarsWithOcc
        , allRelevantVars, allRelevantVarsIgnoring
        , freeIn, freeInIgnoringSorts, isBinderUsed
        , relevantIn, relevantInIgnoringSortAnn
        , Occurrence(..)
        , VarOcc(..)
        , occurrence
        , closed
        , freeVars -- only for testing
        , freeVars'
        , MetaSet
        ) where

    import Prelude hiding (null)

    import Control.Monad.Reader

    import Data.Maybe
    import Data.Monoid ( Monoid, mempty, mappend, mconcat )
    import Data.Semigroup ( Semigroup, (<>), Any(..), All(..) )
    import Data.IntSet (IntSet)
    import qualified Data.IntSet as Set
    import Data.IntMap (IntMap)
    import qualified Data.IntMap as Map
    import Data.Set (Set)
    import Data.Proxy

    import qualified Agda.Benchmarking as Bench

    import Agda.Syntax.Common hiding (Arg, Dom, NamedArg)
    import Agda.Syntax.Internal

    import Agda.TypeChecking.Free.Lazy
      ( Free(..) , FreeEnv(..), initFreeEnv
      , VarOcc(..), topVarOcc, TheVarMap, theVarMap, IgnoreSorts(..), Variable, SingleVar
      , MetaSet, IsVarSet(..), runFreeM
      )
    import qualified Agda.TypeChecking.Free.Lazy as Free

    import Agda.Utils.Null
    import Agda.Utils.Singleton

    type VarSet = IntSet

    -- | Free variables of a term, (disjointly) partitioned into strongly and
    --   and weakly rigid variables, flexible variables and irrelevant variables.
    data FreeVars = FV
      { stronglyRigidVars :: VarSet
        -- ^ Variables under only and at least one inductive constructor(s).
      , unguardedVars     :: VarSet
        -- ^ Variables at top or only under inductive record constructors
        --   λs and Πs.
        --   The purpose of recording these separately is that they
        --   can still become strongly rigid if put under a constructor
        --   whereas weakly rigid ones stay weakly rigid.
      , weaklyRigidVars   :: VarSet
        -- ^ Ordinary rigid variables, e.g., in arguments of variables or functions.
      , flexibleVars      :: IntMap MetaSet
        -- ^ Variables occuring in arguments of metas.
        --   These are only potentially free, depending how the meta variable is instantiated.
        --   The set contains the id's of the meta variables that this variable is an argument to.
      , irrelevantVars    :: VarSet
        -- ^ Variables in irrelevant arguments and under a @DontCare@, i.e.,
        --   in irrelevant positions.
      } deriving (Eq, Show)

    mapSRV, mapUGV, mapWRV, mapIRV
      :: (VarSet -> VarSet) -> FreeVars -> FreeVars
    mapSRV f fv = fv { stronglyRigidVars = f $ stronglyRigidVars fv }
    mapUGV f fv = fv { unguardedVars     = f $ unguardedVars     fv }
    mapWRV f fv = fv { weaklyRigidVars   = f $ weaklyRigidVars   fv }
    mapIRV f fv = fv { irrelevantVars    = f $ irrelevantVars    fv }

    mapFXV :: (IntMap MetaSet -> IntMap MetaSet) -> FreeVars -> FreeVars
    mapFXV f fv = fv { flexibleVars      = f $ flexibleVars      fv }

    -- | Rigid variables: either strongly rigid, unguarded, or weakly rigid.
    rigidVars :: FreeVars -> VarSet
    rigidVars fv = Set.unions
      [ stronglyRigidVars fv
      ,     unguardedVars fv
      ,   weaklyRigidVars fv
      ]

    -- | All but the irrelevant variables.
    relevantVars :: FreeVars -> VarSet
    relevantVars fv = Set.unions [rigidVars fv, Map.keysSet (flexibleVars fv)]

    -- | @allVars fv@ includes irrelevant variables.
    allVars :: FreeVars -> VarSet
    allVars fv = Set.unions [relevantVars fv, irrelevantVars fv]

    data Occurrence
      = NoOccurrence
      | Irrelevantly
      | StronglyRigid     -- ^ Under at least one and only inductive constructors.
      | Unguarded         -- ^ In top position, or only under inductive record constructors.
      | WeaklyRigid       -- ^ In arguments to variables and definitions.
      | Flexible MetaSet  -- ^ In arguments of metas.
      deriving (Eq,Show)

    -- | Compute an occurrence of a single variable in a piece of internal syntax.
    occurrence :: Free a => Nat -> a -> Occurrence
    occurrence x v = occurrenceFV x $ freeVars v

    -- | Extract occurrence of a single variable from computed free variables.
    occurrenceFV :: Nat -> FreeVars -> Occurrence
    occurrenceFV x fv
      | x `Set.member` stronglyRigidVars fv = StronglyRigid
      | x `Set.member` unguardedVars     fv = Unguarded
      | x `Set.member` weaklyRigidVars   fv = WeaklyRigid
      | Just ms <- Map.lookup x (flexibleVars fv) = Flexible ms
      | x `Set.member` irrelevantVars    fv = Irrelevantly
      | otherwise                           = NoOccurrence

    -- | Mark variables as flexible.  Useful when traversing arguments of metas.
    flexible :: MetaSet -> FreeVars -> FreeVars
    flexible ms fv =
        fv { stronglyRigidVars = Set.empty
           , unguardedVars     = Set.empty
           , weaklyRigidVars   = Set.empty
           , flexibleVars      = Map.unionsWith mappend
                                   [ Map.fromSet (const ms) (rigidVars fv)
                                   , fmap (mappend ms) (flexibleVars fv) ]
           }

    -- | Mark rigid variables as non-strongly.  Useful when traversing arguments of variables.
    weakly :: FreeVars -> FreeVars
    weakly fv = fv
      { stronglyRigidVars = Set.empty
      , unguardedVars     = Set.empty
      , weaklyRigidVars   = rigidVars fv
      }

    -- | Mark unguarded variables as strongly rigid.  Useful when traversing arguments of inductive constructors.
    strongly :: FreeVars -> FreeVars
    strongly fv = fv
      { stronglyRigidVars = stronglyRigidVars fv `Set.union` unguardedVars fv
      , unguardedVars     = Set.empty
      }

    -- | What happens to the variables occurring under a constructor?
    underConstructor :: ConHead -> FreeVars -> FreeVars
    underConstructor (ConHead c i fs) =
      case (i,fs) of
        -- Coinductive (record) constructors admit infinite cycles:
        (CoInductive, _)   -> weakly
        -- Inductive data constructors do not admit infinite cycles:
        (Inductive, [])    -> strongly
        -- Inductive record constructors do not admit infinite cycles,
        -- but this cannot be proven inside Agda.
        -- Thus, unification should not prove it either.
        (Inductive, (_:_)) -> id

    -- | Mark all free variables as irrelevant.
    irrelevantly :: FreeVars -> FreeVars
    irrelevantly fv = empty { irrelevantVars = allVars fv }

    -- | Pointwise union.
    union :: FreeVars -> FreeVars -> FreeVars
    union (FV sv1 gv1 rv1 fv1 iv1) (FV sv2 gv2 rv2 fv2 iv2) =
      FV (Set.union sv1 sv2) (Set.union gv1 gv2) (Set.union rv1 rv2) (Map.unionWith mappend fv1 fv2) (Set.union iv1 iv2)

    unions :: [FreeVars] -> FreeVars
    unions = foldr union empty

    instance Null FreeVars where
      empty = FV Set.empty Set.empty Set.empty Map.empty Set.empty
      null (FV a b c d e) = null a && null b && null c && null d && null e

    -- | Free variable sets form a monoid under 'union'.
    instance Semigroup FreeVars where
      (<>) = union

    instance Monoid FreeVars where
      mempty  = empty
      mappend = (<>)
      mconcat = unions

    -- | @delete x fv@ deletes variable @x@ from variable set @fv@.
    delete :: Nat -> FreeVars -> FreeVars
    delete n (FV sv gv rv fv iv) = FV (Set.delete n sv) (Set.delete n gv) (Set.delete n rv) (Map.delete n fv) (Set.delete n iv)

    instance Singleton Variable FreeVars where
      singleton i = mapUGV (Set.insert i) mempty

    instance IsVarSet FreeVars where
      withVarOcc (VarOcc o r) = goOcc o . goRel r
        where
          goOcc o = case o of
            Free.Flexible ms   -> flexible ms
            Free.WeaklyRigid   -> weakly
            Free.Unguarded     -> id
            Free.StronglyRigid -> strongly
          goRel r = case r of
            Relevant   -> id
            NonStrict  -> id    -- we don't track non-strict in FreeVars
            Irrelevant -> irrelevantly

    -- In most cases we don't care about the VarOcc.

    instance IsVarSet VarSet where withVarOcc _ = id
    instance IsVarSet [Int]  where withVarOcc _ = id
    instance IsVarSet Any    where withVarOcc _ = id
    instance IsVarSet All    where withVarOcc _ = id

    newtype VarCounts = VarCounts { varCounts :: IntMap Int }

    instance Semigroup VarCounts where
      VarCounts fv1 <> VarCounts fv2 = VarCounts (Map.unionWith (+) fv1 fv2)

    instance Monoid VarCounts where
      mempty = VarCounts Map.empty
      mappend = (<>)

    instance IsVarSet VarCounts where withVarOcc _ = id

    instance Singleton Variable VarCounts where
      singleton i = VarCounts $ Map.singleton i 1

    -- * Collecting free variables.

    bench :: a -> a
    bench = Bench.billToPure [ Bench.Typing , Bench.Free ]

    -- | Doesn't go inside solved metas, but collects the variables from a
    -- metavariable application @X ts@ as @flexibleVars@.
    {-# SPECIALIZE freeVars :: Free a => a -> FreeVars #-}
    freeVars :: (IsVarSet c, Singleton Variable c, Free a) => a -> c
    freeVars = freeVarsIgnore IgnoreNot

    {-# SPECIALIZE freeVarsIgnore :: Free a => IgnoreSorts -> a -> FreeVars #-}
    freeVarsIgnore :: (IsVarSet c, Singleton Variable c, Free a) =>
                      IgnoreSorts -> a -> c
    freeVarsIgnore = runFree singleton

    -- Specialization to typical monoids
    {-# SPECIALIZE runFree :: Free a => SingleVar Any      -> IgnoreSorts -> a -> Any #-}
    {-# SPECIALIZE runFree :: Free a => SingleVar FreeVars -> IgnoreSorts -> a -> FreeVars #-}
    -- Specialization to Term
    {-# SPECIALIZE runFree :: SingleVar Any      -> IgnoreSorts -> Term -> Any #-}
    {-# SPECIALIZE runFree :: SingleVar FreeVars -> IgnoreSorts -> Term -> FreeVars #-}

    -- | Compute free variables.
    runFree :: (IsVarSet c, Free a) => SingleVar c -> IgnoreSorts -> a -> c
    runFree single i t = -- bench $  -- Benchmarking is expensive (4% on std-lib)
      runFreeM single i (freeVars' t)

    -- | Check if a variable is free, possibly ignoring sorts.
    freeIn' :: Free a => IgnoreSorts -> Nat -> a -> Bool
    freeIn' ig x t = getAny $ runFree (Any . (x ==)) ig t

    {-# SPECIALIZE freeIn :: Nat -> Term -> Bool #-}
    freeIn :: Free a => Nat -> a -> Bool
    freeIn = freeIn' IgnoreNot

    freeInIgnoringSorts :: Free a => Nat -> a -> Bool
    freeInIgnoringSorts = freeIn' IgnoreAll

    freeInIgnoringSortAnn :: Free a => Nat -> a -> Bool
    freeInIgnoringSortAnn = freeIn' IgnoreInAnnotations

    newtype RelevantIn a = RelevantIn {getRelevantIn :: a}
      deriving (Semigroup, Monoid)

    instance IsVarSet a => IsVarSet (RelevantIn a) where
      withVarOcc o x
        | isIrrelevant (varRelevance o) = mempty
        | otherwise = RelevantIn $ withVarOcc o $ getRelevantIn x

    relevantIn' :: Free a => IgnoreSorts -> Nat -> a -> Bool
    relevantIn' ig x t = getAny . getRelevantIn $ runFree (RelevantIn . Any . (x ==)) ig t

    relevantInIgnoringSortAnn :: Free a => Nat -> a -> Bool
    relevantInIgnoringSortAnn = relevantIn' IgnoreInAnnotations

    relevantIn :: Free a => Nat -> a -> Bool
    relevantIn = relevantIn' IgnoreAll

    -- | Is the variable bound by the abstraction actually used?
    isBinderUsed :: Free a => Abs a -> Bool
    isBinderUsed NoAbs{}   = False
    isBinderUsed (Abs _ x) = 0 `freeIn` x

    -- | Is the term entirely closed (no free variables)?
    closed :: Free a => a -> Bool
    closed t = getAll $ runFree (const $ All False) IgnoreNot t

    -- | Collect all free variables.
    allFreeVars :: Free a => a -> VarSet
    allFreeVars = runFree Set.singleton IgnoreNot

    -- | Collect all free variables together with information about their occurrence.
    allFreeVarsWithOcc :: Free a => a -> TheVarMap
    allFreeVarsWithOcc = theVarMap . runFree (singleton . (,topVarOcc)) IgnoreNot

    -- | Collect all relevant free variables, possibly ignoring sorts.
    allRelevantVarsIgnoring :: Free a => IgnoreSorts -> a -> VarSet
    allRelevantVarsIgnoring ig = getRelevantIn . runFree (RelevantIn . Set.singleton) ig

    -- | Collect all relevant free variables, excluding the "unused" ones.
    allRelevantVars :: Free a => a -> VarSet
    allRelevantVars = allRelevantVarsIgnoring IgnoreNot