{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# OPTIONS_GHC -Wunused-imports #-}
#if  __GLASGOW_HASKELL__ > 902
{-# OPTIONS_GHC -fworker-wrapper-cbv #-}
#endif

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

The most general function here is 'freeVarMap', which returns returns occurrence information about
every free variable. It is also a legacy function in the sense that it used to be the only function
for computing free occurrences, and it came with various functions for taking views of the resulting
'VarMap'. You can find these under the "Legacy API" here.

There are also also a bunch of specialized and optimized implementations of traversals here.

If you want to write a new function for computing free occurrence information, your task is
essentially to write a 'ComputeFree' instance and invoke 'freeVars'. These together implement a
generic traversal which has a good chance of being converted to decent code by GHC. You can look at
examples here and also look at the comments at 'ComputeFree'.
-}

module Mikan.TypeChecking.Free
    (
      FlexRig
    , FlexRig'(..)
    , Free(..)
    , IgnoreSorts(..)
    , LensFlexRig(..)
    , VarCounts(..)
    , VarMap
    , VarMap'(..)
    , VarOcc
    , VarOcc'(..)
    , composeFlexRig
    , lookupVarMap
    , mapVarMap

    -- * Legacy API
    , allVars
    , filterVarMap
    , filterVarMapToList
    , flexibleVars
    , freeVarMap
    , isFlexible
    , isStronglyRigid
    , isUnguarded
    , isWeaklyRigid
    , rigidVars
    , stronglyRigidVars
    , unguardedVars

    -- * MetaSet
    , MetaSet(..)
    , foldrMetaSet
    , insertMetaSet
    , metaSetToBlocker

    -- * Specialized traversals
    , allFreeVar
    , anyFreeVar
    , closed
    , flexRigOccurrenceIn
    , freeIn
    , freeInIgnoringSorts
    , freeVarCounts
    , freeVarList
    , freeVarMapIgnoreAnn
    , freeVarSet
    , isBinderUsed
    , setFreeInIgnoring
    , flexRigToBlocker
    ) where

import Prelude hiding (null)
import GHC.Exts (Int(..), Int#, (-#), oneShot)

import Data.Semigroup (Any(..), All(..))
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap

import Mikan.Syntax.Common hiding (Arg, NamedArg)
import Mikan.Syntax.Internal

import Mikan.TypeChecking.Free.Base
import Mikan.TypeChecking.Free.Generic

import Mikan.Utils.VarSet (VarSet)
import Mikan.Utils.VarSet qualified as VarSet
import Mikan.Utils.StrictReader
import Mikan.Utils.StrictFlipEndo
import Mikan.Utils.StrictEndo qualified as NonFlip
import Mikan.Utils.ExpandCase

-- ** All free variables together with information about their occurrence.
--------------------------------------------------------------------------------

data FreeVarMap = FreeVarMap !Int {-# UNPACK #-} !VarOcc

instance LensFlexRig FreeVarMap MetaSet where
  {-# INLINE lensFlexRig #-}
  lensFlexRig :: Lens' FreeVarMap (FlexRig' MetaSet)
lensFlexRig FlexRig' MetaSet -> f (FlexRig' MetaSet)
f (FreeVarMap Int
x VarOcc
occ) = Int -> VarOcc -> FreeVarMap
FreeVarMap Int
x (VarOcc -> FreeVarMap) -> f VarOcc -> f FreeVarMap
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FlexRig' MetaSet -> f (FlexRig' MetaSet)) -> VarOcc -> f VarOcc
forall o a. LensFlexRig o a => Lens' o (FlexRig' a)
Lens' VarOcc (FlexRig' MetaSet)
lensFlexRig FlexRig' MetaSet -> f (FlexRig' MetaSet)
f VarOcc
occ

instance ComputeFree FreeVarMap where
  type Collect FreeVarMap = Endo (IntMap VarOcc)
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> FreeVarMap -> FreeVarMap
underBinders' Int
n (FreeVarMap Int
x VarOcc
occ) = Int -> VarOcc -> FreeVarMap
FreeVarMap (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) VarOcc
occ
  {-# INLINE variable' #-}
  variable' :: Int -> FreeVarMap -> Collect FreeVarMap
variable' Int
x' (FreeVarMap Int
x VarOcc
occ) = (IntMap VarOcc -> IntMap VarOcc) -> Endo (IntMap VarOcc)
forall a. (a -> a) -> Endo a
Endo \IntMap VarOcc
vars ->
    if Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x then IntMap VarOcc
vars
              else (VarOcc -> VarOcc -> VarOcc)
-> Int -> VarOcc -> IntMap VarOcc -> IntMap VarOcc
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith VarOcc -> VarOcc -> VarOcc
forall a. Semigroup a => a -> a -> a
(<>) (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) VarOcc
occ IntMap VarOcc
vars
  underConstructor' :: ConHead -> Elims -> FreeVarMap -> FreeVarMap
underConstructor' = ConHead -> Elims -> FreeVarMap -> FreeVarMap
forall a r.
(Semigroup a, LensFlexRig r a) =>
ConHead -> Elims -> r -> r
defaultUnderConstructor; {-# INLINE underConstructor' #-}
  underFlexRig' :: Maybe (FlexRig' MetaSet -> FreeVarMap -> FreeVarMap)
underFlexRig'     = Maybe (FlexRig' MetaSet -> FreeVarMap -> FreeVarMap)
forall a r.
(Semigroup a, LensFlexRig r a) =>
Maybe (FlexRig' a -> r -> r)
defaultUnderFlexRig;     {-# INLINE underFlexRig' #-}

{-# SPECIALIZE freeVarMap :: Term -> VarMap #-}
{-# SPECIALIZE freeVarMap :: Type -> VarMap #-}
-- | Return information about every free variable.
freeVarMap :: Free t => t -> VarMap
freeVarMap :: forall t. Free t => t -> VarMap
freeVarMap t
t =
  IntMap VarOcc -> VarMap
forall a. TheVarMap' a -> VarMap' a
VarMap (Endo (IntMap VarOcc) -> IntMap VarOcc -> IntMap VarOcc
forall a. Endo a -> a -> a
appEndo (Reader FreeVarMap (Endo (IntMap VarOcc))
-> FreeVarMap -> Endo (IntMap VarOcc)
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeVarMap (Collect FreeVarMap)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> VarOcc -> FreeVarMap
FreeVarMap Int
0 VarOcc
forall a. VarOcc' a
oneVarOcc)) IntMap VarOcc
forall a. Monoid a => a
mempty)

--------------------------------------------------------------------------------

newtype FreeVarMapIgnoreAnn = FreeVarMapIgnoreAnn FreeVarMap

instance LensFlexRig FreeVarMapIgnoreAnn MetaSet where
  {-# INLINE lensFlexRig #-}
  lensFlexRig :: Lens' FreeVarMapIgnoreAnn (FlexRig' MetaSet)
lensFlexRig FlexRig' MetaSet -> f (FlexRig' MetaSet)
f (FreeVarMapIgnoreAnn FreeVarMap
x) = FreeVarMap -> FreeVarMapIgnoreAnn
FreeVarMapIgnoreAnn (FreeVarMap -> FreeVarMapIgnoreAnn)
-> f FreeVarMap -> f FreeVarMapIgnoreAnn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FlexRig' MetaSet -> f (FlexRig' MetaSet))
-> FreeVarMap -> f FreeVarMap
forall o a. LensFlexRig o a => Lens' o (FlexRig' a)
Lens' FreeVarMap (FlexRig' MetaSet)
lensFlexRig FlexRig' MetaSet -> f (FlexRig' MetaSet)
f FreeVarMap
x

instance ComputeFree FreeVarMapIgnoreAnn where
  type Collect FreeVarMapIgnoreAnn = Endo (IntMap VarOcc)
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> FreeVarMapIgnoreAnn -> FreeVarMapIgnoreAnn
underBinders' Int
n (FreeVarMapIgnoreAnn FreeVarMap
x) = FreeVarMap -> FreeVarMapIgnoreAnn
FreeVarMapIgnoreAnn (Int -> FreeVarMap -> FreeVarMap
forall r. ComputeFree r => Int -> r -> r
underBinders' Int
n FreeVarMap
x)
  {-# INLINE variable' #-}
  variable' :: Int -> FreeVarMapIgnoreAnn -> Collect FreeVarMapIgnoreAnn
variable' Int
x' (FreeVarMapIgnoreAnn FreeVarMap
x) = Int -> FreeVarMap -> Collect FreeVarMap
forall r. ComputeFree r => Int -> r -> Collect r
variable' Int
x' FreeVarMap
x
  ignoreSorts' :: IgnoreSorts
ignoreSorts'      = IgnoreSorts
IgnoreInAnnotations;     {-# INLINE ignoreSorts' #-}
  underConstructor' :: ConHead -> Elims -> FreeVarMapIgnoreAnn -> FreeVarMapIgnoreAnn
underConstructor' = ConHead -> Elims -> FreeVarMapIgnoreAnn -> FreeVarMapIgnoreAnn
forall a r.
(Semigroup a, LensFlexRig r a) =>
ConHead -> Elims -> r -> r
defaultUnderConstructor; {-# INLINE underConstructor' #-}
  underFlexRig' :: Maybe
  (FlexRig' MetaSet -> FreeVarMapIgnoreAnn -> FreeVarMapIgnoreAnn)
underFlexRig'     = Maybe
  (FlexRig' MetaSet -> FreeVarMapIgnoreAnn -> FreeVarMapIgnoreAnn)
forall a r.
(Semigroup a, LensFlexRig r a) =>
Maybe (FlexRig' a -> r -> r)
defaultUnderFlexRig;     {-# INLINE underFlexRig' #-}

{-# SPECIALIZE freeVarMapIgnoreAnn :: Term -> VarMap #-}
{-# SPECIALIZE freeVarMapIgnoreAnn :: (Term,Type) -> VarMap #-}
-- | Return information about every free variable, but ignore occurrences
--   in sorts of types.
freeVarMapIgnoreAnn :: Free t => t -> VarMap
freeVarMapIgnoreAnn :: forall t. Free t => t -> VarMap
freeVarMapIgnoreAnn t
t =
  IntMap VarOcc -> VarMap
forall a. TheVarMap' a -> VarMap' a
VarMap (Endo (IntMap VarOcc) -> IntMap VarOcc -> IntMap VarOcc
forall a. Endo a -> a -> a
appEndo (Reader FreeVarMapIgnoreAnn (Endo (IntMap VarOcc))
-> FreeVarMapIgnoreAnn -> Endo (IntMap VarOcc)
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeVarMapIgnoreAnn (Collect FreeVarMapIgnoreAnn)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t)
                  (FreeVarMap -> FreeVarMapIgnoreAnn
FreeVarMapIgnoreAnn (Int -> VarOcc -> FreeVarMap
FreeVarMap Int
0 VarOcc
forall a. VarOcc' a
oneVarOcc))) IntMap VarOcc
forall a. Monoid a => a
mempty)

--------------------------------------------------------------------------------

newtype FreeVarCounts = FreeVarCounts Int

instance ComputeFree FreeVarCounts where
  type Collect FreeVarCounts = Endo (IntMap Int)
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> FreeVarCounts -> FreeVarCounts
underBinders' Int
n (FreeVarCounts Int
x) = Int -> FreeVarCounts
FreeVarCounts (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
  {-# INLINE variable' #-}
  variable' :: Int -> FreeVarCounts -> Collect FreeVarCounts
variable' Int
x' (FreeVarCounts Int
x) = (IntMap Int -> IntMap Int) -> Endo (IntMap Int)
forall a. (a -> a) -> Endo a
Endo \IntMap Int
counts ->
    if Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x then IntMap Int
counts else (Int -> Int -> Int) -> Int -> Int -> IntMap Int -> IntMap Int
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) Int
1 IntMap Int
counts

{-# SPECIALIZE freeVarCounts :: Term -> VarCounts #-}
-- | Return how many times each free variable occurs.
freeVarCounts :: Free t => t -> VarCounts
freeVarCounts :: forall t. Free t => t -> VarCounts
freeVarCounts t
t =
  IntMap Int -> VarCounts
VarCounts (Endo (IntMap Int) -> IntMap Int -> IntMap Int
forall a. Endo a -> a -> a
appEndo (Reader FreeVarCounts (Endo (IntMap Int))
-> FreeVarCounts -> Endo (IntMap Int)
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeVarCounts (Collect FreeVarCounts)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> FreeVarCounts
FreeVarCounts Int
0)) IntMap Int
forall a. Monoid a => a
mempty)

-- ** Testing that a predicate holds on any free variable
--------------------------------------------------------------------------------

data AnyFreeVar = AnyFreeVar !Int !(Int# -> Bool)

instance ComputeFree AnyFreeVar where
  type Collect AnyFreeVar = Any
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> AnyFreeVar -> AnyFreeVar
underBinders' Int
n (AnyFreeVar Int
x Int# -> Bool
f) = Int -> (Int# -> Bool) -> AnyFreeVar
AnyFreeVar (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) Int# -> Bool
f
  {-# INLINE variable' #-}
  variable' :: Int -> AnyFreeVar -> Collect AnyFreeVar
variable' (I# Int#
x') (AnyFreeVar (I# Int#
x) Int# -> Bool
f) = if Int# -> Int
I# Int#
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int# -> Int
I# Int#
x then Any
Collect AnyFreeVar
forall a. Monoid a => a
mempty else Bool -> Any
Any (Int# -> Bool
f (Int#
x' Int# -> Int# -> Int#
-# Int#
x))

{-# SPECIALIZE anyFreeVar# :: (Int# -> Bool) -> Dom' Term Type -> Bool #-}
anyFreeVar# :: Free t => (Int# -> Bool) -> t -> Bool
anyFreeVar# :: forall t. Free t => (Int# -> Bool) -> t -> Bool
anyFreeVar# Int# -> Bool
f t
t = Any -> Bool
getAny (Reader AnyFreeVar Any -> AnyFreeVar -> Any
forall r a. Reader r a -> r -> a
runReader (t -> Reader AnyFreeVar (Collect AnyFreeVar)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> (Int# -> Bool) -> AnyFreeVar
AnyFreeVar Int
0 Int# -> Bool
f))

{-# INLINE anyFreeVar #-}
-- | Compute the disjunction of a predicate on free variables.
anyFreeVar :: Free t => (Int -> Bool) -> t -> Bool
anyFreeVar :: forall t. Free t => (Int -> Bool) -> t -> Bool
anyFreeVar Int -> Bool
f = (Int# -> Bool) -> t -> Bool
forall t. Free t => (Int# -> Bool) -> t -> Bool
anyFreeVar# (\Int#
n -> Int -> Bool
f (Int# -> Int
I# Int#
n))

{-# SPECIALIZE allFreeVar# :: (Int# -> Bool) -> Dom' Term Type -> Bool #-}
allFreeVar# :: Free t => (Int# -> Bool) -> t -> Bool
allFreeVar# :: forall t. Free t => (Int# -> Bool) -> t -> Bool
allFreeVar# Int# -> Bool
f t
t = Bool -> Bool
not (Any -> Bool
getAny (Reader AnyFreeVar Any -> AnyFreeVar -> Any
forall r a. Reader r a -> r -> a
runReader (t -> Reader AnyFreeVar (Collect AnyFreeVar)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> (Int# -> Bool) -> AnyFreeVar
AnyFreeVar Int
0 (\Int#
n -> Bool -> Bool
not (Int# -> Bool
f Int#
n)))))

{-# INLINE allFreeVar #-}
-- | Compute the conjunction of a predicate on free variables.
allFreeVar :: Free t => (Int -> Bool) -> t -> Bool
allFreeVar :: forall t. Free t => (Int -> Bool) -> t -> Bool
allFreeVar Int -> Bool
f = (Int# -> Bool) -> t -> Bool
forall t. Free t => (Int# -> Bool) -> t -> Bool
allFreeVar# (\Int#
n -> Int -> Bool
f (Int# -> Int
I# Int#
n))

-- ** Flex-rigid occurrence for a single variable
--------------------------------------------------------------------------------

data SingleFR = SingleFR !Int !FlexRig

instance LensFlexRig SingleFR MetaSet where
  {-# INLINE lensFlexRig #-}
  lensFlexRig :: Lens' SingleFR (FlexRig' MetaSet)
lensFlexRig FlexRig' MetaSet -> f (FlexRig' MetaSet)
f (SingleFR Int
x FlexRig' MetaSet
fr) = Int -> FlexRig' MetaSet -> SingleFR
SingleFR Int
x (FlexRig' MetaSet -> SingleFR)
-> f (FlexRig' MetaSet) -> f SingleFR
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FlexRig' MetaSet -> f (FlexRig' MetaSet)
f FlexRig' MetaSet
fr

data CollectSingleFR = CSFRNothing | CSFRJust !FlexRig

instance Semigroup CollectSingleFR where
  CSFRJust FlexRig' MetaSet
fr <> :: CollectSingleFR -> CollectSingleFR -> CollectSingleFR
<> CSFRJust FlexRig' MetaSet
fr' = FlexRig' MetaSet -> CollectSingleFR
CSFRJust (FlexRig' MetaSet -> FlexRig' MetaSet -> FlexRig' MetaSet
forall a. Semigroup a => FlexRig' a -> FlexRig' a -> FlexRig' a
addFlexRig FlexRig' MetaSet
fr FlexRig' MetaSet
fr')
  CollectSingleFR
CSFRNothing <> CollectSingleFR
s            = CollectSingleFR
s
  CollectSingleFR
s           <> CollectSingleFR
CSFRNothing  = CollectSingleFR
s

instance ComputeFree SingleFR where
  type Collect SingleFR = Endo CollectSingleFR
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> SingleFR -> SingleFR
underBinders' Int
n (SingleFR Int
x FlexRig' MetaSet
fr) = Int -> FlexRig' MetaSet -> SingleFR
SingleFR (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) FlexRig' MetaSet
fr
  {-# INLINE variable' #-}
  variable' :: Int -> SingleFR -> Collect SingleFR
variable' Int
x' (SingleFR Int
x FlexRig' MetaSet
fr) = (CollectSingleFR -> CollectSingleFR) -> Endo CollectSingleFR
forall a. (a -> a) -> Endo a
Endo \CollectSingleFR
acc -> if Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
x' then CollectSingleFR
acc CollectSingleFR -> CollectSingleFR -> CollectSingleFR
forall a. Semigroup a => a -> a -> a
<> FlexRig' MetaSet -> CollectSingleFR
CSFRJust FlexRig' MetaSet
fr else CollectSingleFR
acc
  underConstructor' :: ConHead -> Elims -> SingleFR -> SingleFR
underConstructor' = ConHead -> Elims -> SingleFR -> SingleFR
forall a r.
(Semigroup a, LensFlexRig r a) =>
ConHead -> Elims -> r -> r
defaultUnderConstructor; {-# INLINE underConstructor' #-}
  underFlexRig' :: Maybe (FlexRig' MetaSet -> SingleFR -> SingleFR)
underFlexRig'     = Maybe (FlexRig' MetaSet -> SingleFR -> SingleFR)
forall a r.
(Semigroup a, LensFlexRig r a) =>
Maybe (FlexRig' a -> r -> r)
defaultUnderFlexRig; {-# INLINE underFlexRig' #-}

{-# SPECIALIZE flexRigOccurrenceIn :: Nat -> Term -> Maybe FlexRig #-}
{-# SPECIALIZE flexRigOccurrenceIn :: Nat -> Sort' Term -> Maybe FlexRig #-}
-- | Compute 'FlexRig' for a single variable.
flexRigOccurrenceIn :: Free a => Nat -> a -> Maybe FlexRig
flexRigOccurrenceIn :: forall a. Free a => Int -> a -> Maybe (FlexRig' MetaSet)
flexRigOccurrenceIn Int
x a
a = case Endo CollectSingleFR -> CollectSingleFR -> CollectSingleFR
forall a. Endo a -> a -> a
appEndo (Reader SingleFR (Endo CollectSingleFR)
-> SingleFR -> Endo CollectSingleFR
forall r a. Reader r a -> r -> a
runReader (a -> Reader SingleFR (Collect SingleFR)
forall r. ComputeFree r => a -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars a
a) (Int -> FlexRig' MetaSet -> SingleFR
SingleFR Int
x FlexRig' MetaSet
forall a. FlexRig' a
Unguarded)) CollectSingleFR
CSFRNothing of
  CollectSingleFR
CSFRNothing -> Maybe (FlexRig' MetaSet)
forall a. Maybe a
Nothing
  CSFRJust FlexRig' MetaSet
fr -> FlexRig' MetaSet -> Maybe (FlexRig' MetaSet)
forall a. a -> Maybe a
Just FlexRig' MetaSet
fr

-- ** Plain free occurrence
--------------------------------------------------------------------------------

newtype FreeIn = FreeIn Int

instance ComputeFree FreeIn where
  type Collect FreeIn = Any
  underBinders' :: Int -> FreeIn -> FreeIn
underBinders' Int
n (FreeIn Int
x) = Int -> FreeIn
FreeIn (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x); {-# INLINE underBinders' #-}
  variable' :: Int -> FreeIn -> Collect FreeIn
variable' Int
x' (FreeIn Int
x) = Bool -> Any
Any (Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
x'); {-# INLINE variable' #-}

{-# SPECIALIZE freeIn :: Nat -> Term -> Bool #-}
{-# SPECIALIZE freeIn :: Nat -> Type -> Bool #-}
-- | Check if single variable occurs freely.
freeIn :: Free t => Nat -> t -> Bool
freeIn :: forall t. Free t => Int -> t -> Bool
freeIn Int
x t
a = Any -> Bool
getAny (Reader FreeIn Any -> FreeIn -> Any
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeIn (Collect FreeIn)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
a) (Int -> FreeIn
FreeIn Int
x))

newtype FreeInIgnoringSorts = FreeInIgnoringSorts Int

instance ComputeFree FreeInIgnoringSorts where
  type Collect FreeInIgnoringSorts = Any
  underBinders' :: Int -> FreeInIgnoringSorts -> FreeInIgnoringSorts
underBinders' Int
n (FreeInIgnoringSorts Int
x) = Int -> FreeInIgnoringSorts
FreeInIgnoringSorts (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x); {-# INLINE underBinders' #-}
  variable' :: Int -> FreeInIgnoringSorts -> Collect FreeInIgnoringSorts
variable' Int
x' (FreeInIgnoringSorts Int
x) = Bool -> Any
Any (Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
x'); {-# INLINE variable' #-}
  ignoreSorts' :: IgnoreSorts
ignoreSorts' = IgnoreSorts
IgnoreAll; {-# INLINE ignoreSorts' #-}

{-# SPECIALIZE freeInIgnoringSorts :: Nat -> Term -> Bool #-}
{-# SPECIALIZE freeInIgnoringSorts :: Nat -> Type -> Bool #-}
-- | Check if a single variable occurs freely outside of sorts of types.
freeInIgnoringSorts :: Free a => Nat -> a -> Bool
freeInIgnoringSorts :: forall t. Free t => Int -> t -> Bool
freeInIgnoringSorts Int
x a
a = Any -> Bool
getAny (Reader FreeInIgnoringSorts Any -> FreeInIgnoringSorts -> Any
forall r a. Reader r a -> r -> a
runReader (a -> Reader FreeInIgnoringSorts (Collect FreeInIgnoringSorts)
forall r. ComputeFree r => a -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars a
a) (Int -> FreeInIgnoringSorts
FreeInIgnoringSorts Int
x))

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

-- ** Closed objects
---------------------------------------------------------------------------

newtype Closed = Closed Int -- ^ Local scope size.

instance ComputeFree Closed where
  type Collect Closed = All
  underBinders' :: Int -> Closed -> Closed
underBinders' Int
n (Closed Int
x) = Int -> Closed
Closed (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x); {-# INLINE underBinders' #-}
  variable' :: Int -> Closed -> Collect Closed
variable' Int
x' (Closed Int
x) = Bool -> All
All (Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x); {-# INLINE variable' #-}

{-# SPECIALIZE closed :: Term -> Bool #-}
{-# SPECIALIZE closed :: Dom' Term Type -> Bool #-}
-- | Is a term closed?
closed :: Free t => t -> Bool
closed :: forall t. Free t => t -> Bool
closed t
t = All -> Bool
getAll (Reader Closed All -> Closed -> All
forall r a. Reader r a -> r -> a
runReader (t -> Reader Closed (Collect Closed)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> Closed
Closed Int
0))

-- ** Test free occurrence for a set of variables
--------------------------------------------------------------------------------

data SetFreeInIgnoring = SetFreeInIgnoring !Int

newtype CollectSFII = CollectSFII {CollectSFII -> VarSet -> VarSet
appCSFII :: VarSet -> VarSet}

instance Semigroup CollectSFII where
  {-# INLINE (<>) #-}
  CollectSFII VarSet -> VarSet
f <> :: CollectSFII -> CollectSFII -> CollectSFII
<> CollectSFII VarSet -> VarSet
g = (VarSet -> VarSet) -> CollectSFII
CollectSFII ((VarSet -> VarSet) -> VarSet -> VarSet
forall a b. (a -> b) -> a -> b
oneShot \VarSet
xs -> case VarSet -> VarSet
f VarSet
xs of
    VarSet
xs -> if VarSet -> Bool
VarSet.null VarSet
xs then VarSet
xs else VarSet -> VarSet
g VarSet
xs)

instance Monoid CollectSFII where
  {-# INLINE mempty #-}
  mempty :: CollectSFII
mempty = (VarSet -> VarSet) -> CollectSFII
CollectSFII \VarSet
xs -> VarSet
xs

instance ExpandCase LiftedRep CollectSFII where
  type Result LiftedRep CollectSFII = VarSet
  {-# INLINE expand #-}
  expand :: ((CollectSFII -> Result LiftedRep CollectSFII)
 -> Result LiftedRep CollectSFII)
-> CollectSFII
expand (CollectSFII -> Result LiftedRep CollectSFII)
-> Result LiftedRep CollectSFII
k = (VarSet -> VarSet) -> CollectSFII
CollectSFII ((VarSet -> VarSet) -> VarSet -> VarSet
forall a b. (a -> b) -> a -> b
oneShot \VarSet
a -> (CollectSFII -> Result LiftedRep CollectSFII)
-> Result LiftedRep CollectSFII
k ((CollectSFII -> VarSet) -> CollectSFII -> VarSet
forall a b. (a -> b) -> a -> b
oneShot \CollectSFII
act -> CollectSFII -> VarSet -> VarSet
appCSFII CollectSFII
act VarSet
a))

instance ComputeFree SetFreeInIgnoring where
  type Collect SetFreeInIgnoring = CollectSFII
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> SetFreeInIgnoring -> SetFreeInIgnoring
underBinders' Int
n (SetFreeInIgnoring Int
x) = Int -> SetFreeInIgnoring
SetFreeInIgnoring (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
  {-# INLINE variable' #-}
  variable' :: Int -> SetFreeInIgnoring -> Collect SetFreeInIgnoring
variable' Int
x' (SetFreeInIgnoring Int
x) = (VarSet -> VarSet) -> CollectSFII
CollectSFII \VarSet
xs ->
    if Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x
      then Int -> VarSet -> VarSet
VarSet.delete (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) VarSet
xs
      else VarSet
xs
  ignoreSorts' :: IgnoreSorts
ignoreSorts' = IgnoreSorts
IgnoreInAnnotations; {-# INLINE ignoreSorts' #-}

{-# SPECIALIZE NOINLINE setFreeInIgnoring :: VarSet -> Dom Type -> VarSet #-}
{-# SPECIALIZE NOINLINE setFreeInIgnoring :: VarSet -> Type -> VarSet #-}
-- | Test for a set of variables whether each variable occurs outside of sort annotations.
--   Return the set of variables that __don't__ occur.
setFreeInIgnoring :: Free t => VarSet -> t -> VarSet
setFreeInIgnoring :: forall t. Free t => VarSet -> t -> VarSet
setFreeInIgnoring VarSet
xs t
t
  | VarSet -> Bool
VarSet.null VarSet
xs = VarSet
xs
  | Bool
otherwise      = CollectSFII -> VarSet -> VarSet
appCSFII (Reader SetFreeInIgnoring CollectSFII
-> SetFreeInIgnoring -> CollectSFII
forall r a. Reader r a -> r -> a
runReader (t -> Reader SetFreeInIgnoring (Collect SetFreeInIgnoring)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> SetFreeInIgnoring
SetFreeInIgnoring Int
0)) VarSet
xs


-- ** Collect free variables
--------------------------------------------------------------------------------

newtype FreeVarSet = FreeVarSet Int

instance ComputeFree FreeVarSet where
  type Collect FreeVarSet = Endo VarSet
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> FreeVarSet -> FreeVarSet
underBinders' Int
n (FreeVarSet Int
x) = Int -> FreeVarSet
FreeVarSet (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
  {-# INLINE variable' #-}
  variable' :: Int -> FreeVarSet -> Collect FreeVarSet
variable' Int
x' (FreeVarSet Int
x) = (VarSet -> VarSet) -> Endo VarSet
forall a. (a -> a) -> Endo a
Endo \VarSet
xs -> if Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x then VarSet
xs else Int -> VarSet -> VarSet
VarSet.insert (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) VarSet
xs

{-# SPECIALIZE freeVarSet :: Term -> VarSet #-}
{-# SPECIALIZE freeVarSet :: Dom' Term Type -> VarSet #-}
-- | Compute the set of free variables.
freeVarSet :: Free t => t -> VarSet
freeVarSet :: forall t. Free t => t -> VarSet
freeVarSet t
t = Endo VarSet -> VarSet -> VarSet
forall a. Endo a -> a -> a
appEndo (Reader FreeVarSet (Endo VarSet) -> FreeVarSet -> Endo VarSet
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeVarSet (Collect FreeVarSet)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> FreeVarSet
FreeVarSet Int
0)) VarSet
forall a. Monoid a => a
mempty

newtype FreeVarList = FreeVarList Int

instance ComputeFree FreeVarList where
  -- Need the non-flipped endo to get left-right element order
  type Collect FreeVarList = NonFlip.Endo [Int]
  {-# INLINE underBinders' #-}
  underBinders' :: Int -> FreeVarList -> FreeVarList
underBinders' Int
n (FreeVarList Int
x) = Int -> FreeVarList
FreeVarList (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x)
  {-# INLINE variable' #-}
  variable' :: Int -> FreeVarList -> Collect FreeVarList
variable' Int
x' (FreeVarList Int
x) = ([Int] -> [Int]) -> Endo [Int]
forall a. (a -> a) -> Endo a
NonFlip.Endo \[Int]
xs ->
    if Int
x' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x then [Int]
xs else let !v :: Int
v = Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x in Int
v Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
xs

{-# SPECIALIZE freeVarList :: Term -> [Int] #-}
{-# SPECIALIZE freeVarList :: Maybe Term -> [Int] #-}
-- | Compute a (possibly non-unique) list of free variables, in preorder traversal order.
freeVarList :: Free t => t -> [Int]
freeVarList :: forall t. Free t => t -> [Int]
freeVarList t
t = Endo [Int] -> [Int] -> [Int]
forall a. Endo a -> a -> a
NonFlip.appEndo (Reader FreeVarList (Endo [Int]) -> FreeVarList -> Endo [Int]
forall r a. Reader r a -> r -> a
runReader (t -> Reader FreeVarList (Collect FreeVarList)
forall r. ComputeFree r => t -> Reader r (Collect r)
forall t r. (Free t, ComputeFree r) => t -> Reader r (Collect r)
freeVars t
t) (Int -> FreeVarList
FreeVarList Int
0)) []

-- ** Backwards-compatible interface, for extracting information from a 'VarMap'.
---------------------------------------------------------------------------

{-# INLINE filterVarMap #-}
filterVarMap :: (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap :: (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap VarOcc -> Bool
f = \(VarMap IntMap VarOcc
m) ->
  (VarSet -> Int -> VarOcc -> VarSet)
-> VarSet -> IntMap VarOcc -> VarSet
forall a b. (a -> Int -> b -> a) -> a -> IntMap b -> a
IntMap.foldlWithKey'
  (\VarSet
acc Int
k VarOcc
v -> if VarOcc -> Bool
f VarOcc
v then Int -> VarSet -> VarSet
VarSet.insert Int
k VarSet
acc else VarSet
acc)
  VarSet
forall a. Monoid a => a
mempty IntMap VarOcc
m

{-# INLINE filterVarMapToList #-}
filterVarMapToList :: (VarOcc -> Bool) -> VarMap -> [Int]
filterVarMapToList :: (VarOcc -> Bool) -> VarMap -> [Int]
filterVarMapToList VarOcc -> Bool
f = \(VarMap IntMap VarOcc
m) ->
  (Int -> VarOcc -> [Int] -> [Int])
-> [Int] -> IntMap VarOcc -> [Int]
forall a b. (Int -> a -> b -> b) -> b -> IntMap a -> b
IntMap.foldrWithKey'
  (\Int
k VarOcc
v [Int]
acc -> if VarOcc -> Bool
f VarOcc
v then Int
kInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
acc else [Int]
acc)
  [] IntMap VarOcc
m

-- | Variables under only and at least one inductive constructor(s).
stronglyRigidVars :: VarMap -> VarSet
stronglyRigidVars :: VarMap -> VarSet
stronglyRigidVars = (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap ((VarOcc -> Bool) -> VarMap -> VarSet)
-> (VarOcc -> Bool) -> VarMap -> VarSet
forall a b. (a -> b) -> a -> b
$ \case
  VarOcc FlexRig' MetaSet
StronglyRigid -> Bool
True
  VarOcc
_                    -> Bool
False

-- | 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.
unguardedVars :: VarMap -> VarSet
unguardedVars :: VarMap -> VarSet
unguardedVars = (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap ((VarOcc -> Bool) -> VarMap -> VarSet)
-> (VarOcc -> Bool) -> VarMap -> VarSet
forall a b. (a -> b) -> a -> b
$ \case
  VarOcc FlexRig' MetaSet
Unguarded -> Bool
True
  VarOcc
_                -> Bool
False

-- | Rigid variables: either strongly rigid, unguarded, or weakly rigid.
rigidVars :: VarMap -> VarSet
rigidVars :: VarMap -> VarSet
rigidVars = (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap ((VarOcc -> Bool) -> VarMap -> VarSet)
-> (VarOcc -> Bool) -> VarMap -> VarSet
forall a b. (a -> b) -> a -> b
$ \case
  VarOcc FlexRig' MetaSet
o -> case FlexRig' MetaSet
o of
    FlexRig' MetaSet
WeaklyRigid   -> Bool
True
    FlexRig' MetaSet
Unguarded     -> Bool
True
    FlexRig' MetaSet
StronglyRigid -> Bool
True
    FlexRig' MetaSet
_             -> Bool
False

-- | 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.
flexibleVars :: VarMap -> IntMap MetaSet
flexibleVars :: VarMap -> IntMap MetaSet
flexibleVars (VarMap IntMap VarOcc
m) = ((VarOcc -> Maybe MetaSet) -> IntMap VarOcc -> IntMap MetaSet)
-> IntMap VarOcc -> (VarOcc -> Maybe MetaSet) -> IntMap MetaSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip (VarOcc -> Maybe MetaSet) -> IntMap VarOcc -> IntMap MetaSet
forall a b. (a -> Maybe b) -> IntMap a -> IntMap b
IntMap.mapMaybe IntMap VarOcc
m ((VarOcc -> Maybe MetaSet) -> IntMap MetaSet)
-> (VarOcc -> Maybe MetaSet) -> IntMap MetaSet
forall a b. (a -> b) -> a -> b
$ \case
 VarOcc (Flexible MetaSet
ms) -> MetaSet -> Maybe MetaSet
forall a. a -> Maybe a
Just MetaSet
ms
 VarOcc
_                    -> Maybe MetaSet
forall a. Maybe a
Nothing

allVars :: VarMap -> VarSet
allVars :: VarMap -> VarSet
allVars = (VarOcc -> Bool) -> VarMap -> VarSet
filterVarMap (\VarOcc
_ -> Bool
True)