module Agda.Compiler.Treeless.Unused
  ( usedArguments
  , stripUnusedArguments
  ) where

import Data.Maybe

import qualified Data.Set as Set
  -- Andreas, 2021-02-10
  -- TODO: Replace Set by IntSet.
  -- However, this has to wait until we can comfortably bump to
  -- containers-0.6.3.1, which is the first to contain IntSet.mapMonotonic.
  -- Currently, such a constraints gets us into cabal hell.
  -- GHC 8.10 is still shipped with 0.6.2.1, so we either have to wait
  -- until we drop GHC 8 or until we adopt v2-cabal.

import Agda.Syntax.Treeless
import Agda.TypeChecking.Monad
import Agda.TypeChecking.Substitute

import Agda.Compiler.Treeless.Pretty () -- instance only

import Agda.Utils.Function ( iterateUntilM )
import Agda.Utils.List     ( downFrom )
import Agda.Utils.Pretty   ( prettyShow )

usedArguments :: QName -> TTerm -> TCM [ArgUsage]
usedArguments :: QName -> TTerm -> TCM [ArgUsage]
usedArguments QName
q TTerm
t = QName -> TTerm -> [ArgUsage] -> TCM [ArgUsage]
computeUnused QName
q TTerm
b (Int -> ArgUsage -> [ArgUsage]
forall a. Int -> a -> [a]
replicate Int
n ArgUsage
ArgUnused)
  where (Int
n, TTerm
b) = TTerm -> (Int, TTerm)
tLamView TTerm
t

-- | Saturation algorithm, starting with all unused arguments
--   and adding usages until fixed-point has been reached.

computeUnused :: QName -> TTerm -> [ArgUsage] -> TCM [ArgUsage]
computeUnused :: QName -> TTerm -> [ArgUsage] -> TCM [ArgUsage]
computeUnused QName
q TTerm
t = ([ArgUsage] -> [ArgUsage] -> Bool)
-> ([ArgUsage] -> TCM [ArgUsage]) -> [ArgUsage] -> TCM [ArgUsage]
forall (m :: * -> *) a.
Monad m =>
(a -> a -> Bool) -> (a -> m a) -> a -> m a
iterateUntilM [ArgUsage] -> [ArgUsage] -> Bool
forall a. Eq a => a -> a -> Bool
(==) (([ArgUsage] -> TCM [ArgUsage]) -> [ArgUsage] -> TCM [ArgUsage])
-> ([ArgUsage] -> TCM [ArgUsage]) -> [ArgUsage] -> TCM [ArgUsage]
forall a b. (a -> b) -> a -> b
$ \ [ArgUsage]
used -> do

  [Char] -> Int -> [Char] -> TCMT IO ()
forall (m :: * -> *).
MonadDebug m =>
[Char] -> Int -> [Char] -> m ()
reportSLn [Char]
"treeless.opt.unused" Int
50 ([Char] -> TCMT IO ()) -> [Char] -> TCMT IO ()
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ [Char]
"Unused approximation for ", QName -> [Char]
forall a. Pretty a => a -> [Char]
prettyShow QName
q, [Char]
": "
    , [[Char]] -> [Char]
unwords [ if ArgUsage
u ArgUsage -> ArgUsage -> Bool
forall a. Eq a => a -> a -> Bool
== ArgUsage
ArgUsed then [Char
x] else [Char]
"_" | (Char
x, ArgUsage
u) <- [Char] -> [ArgUsage] -> [(Char, ArgUsage)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Char
'a'..] [ArgUsage]
used ]
    ]
  -- Update usage information q to so far "best" value.
  QName -> [ArgUsage] -> TCMT IO ()
setCompiledArgUse QName
q [ArgUsage]
used

  -- The new usage information is the free variables of @t@,
  -- computed under the current usage assumptions of the functions it calls.
  Set Int
fv <- TTerm -> TCMT IO (Set Int)
go TTerm
t
  [ArgUsage] -> TCM [ArgUsage]
forall a. a -> TCMT IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([ArgUsage] -> TCM [ArgUsage]) -> [ArgUsage] -> TCM [ArgUsage]
forall a b. (a -> b) -> a -> b
$ [ if Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
i Set Int
fv then ArgUsage
ArgUsed else ArgUsage
ArgUnused
           | Int
i <- Int -> [Int]
forall a. Integral a => a -> [a]
downFrom ([ArgUsage] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ArgUsage]
used)
           ]
  where
    go :: TTerm -> TCMT IO (Set Int)
go = \case
      TVar Int
x    -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set Int -> TCMT IO (Set Int)) -> Set Int -> TCMT IO (Set Int)
forall a b. (a -> b) -> a -> b
$ Int -> Set Int
forall a. a -> Set a
Set.singleton Int
x
      TPrim{}   -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TDef{}    -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TLit{}    -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TCon{}    -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty

      TApp (TDef QName
f) Args
ts -> do
        [ArgUsage]
used <- [ArgUsage] -> Maybe [ArgUsage] -> [ArgUsage]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [ArgUsage] -> [ArgUsage])
-> TCMT IO (Maybe [ArgUsage]) -> TCM [ArgUsage]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> TCMT IO (Maybe [ArgUsage])
forall (m :: * -> *).
HasConstInfo m =>
QName -> m (Maybe [ArgUsage])
getCompiledArgUse QName
f
        [Set Int] -> Set Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Int] -> Set Int) -> TCMT IO [Set Int] -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TCMT IO (Set Int)] -> TCMT IO [Set Int]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [ TTerm -> TCMT IO (Set Int)
go TTerm
t | (TTerm
t, ArgUsage
ArgUsed) <- Args -> [ArgUsage] -> [(TTerm, ArgUsage)]
forall a b. [a] -> [b] -> [(a, b)]
zip Args
ts ([ArgUsage] -> [(TTerm, ArgUsage)])
-> [ArgUsage] -> [(TTerm, ArgUsage)]
forall a b. (a -> b) -> a -> b
$ [ArgUsage]
used [ArgUsage] -> [ArgUsage] -> [ArgUsage]
forall a. [a] -> [a] -> [a]
++ ArgUsage -> [ArgUsage]
forall a. a -> [a]
repeat ArgUsage
ArgUsed ]

      TApp TTerm
f Args
ts -> [Set Int] -> Set Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Int] -> Set Int) -> TCMT IO [Set Int] -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> TCMT IO (Set Int)) -> Args -> TCMT IO [Set Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TTerm -> TCMT IO (Set Int)
go (TTerm
f TTerm -> Args -> Args
forall a. a -> [a] -> [a]
: Args
ts)
      TLam TTerm
b    -> Set Int -> Set Int
underBinder (Set Int -> Set Int) -> TCMT IO (Set Int) -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCMT IO (Set Int)
go TTerm
b
      TLet TTerm
e TTerm
b  -> do
        Set Int
uses <- TTerm -> TCMT IO (Set Int)
go TTerm
b
        if | Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
0 Set Int
uses -> Set Int -> Set Int -> Set Int
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Set Int -> Set Int
underBinder Set Int
uses) (Set Int -> Set Int) -> TCMT IO (Set Int) -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCMT IO (Set Int)
go TTerm
e
           | Bool
otherwise         -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Set Int -> Set Int
underBinder Set Int
uses)
      TCase Int
x CaseInfo
_ TTerm
d [TAlt]
bs -> Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
x (Set Int -> Set Int)
-> ([Set Int] -> Set Int) -> [Set Int] -> Set Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Set Int] -> Set Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Int] -> Set Int) -> TCMT IO [Set Int] -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((:) (Set Int -> [Set Int] -> [Set Int])
-> TCMT IO (Set Int) -> TCMT IO ([Set Int] -> [Set Int])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCMT IO (Set Int)
go TTerm
d TCMT IO ([Set Int] -> [Set Int])
-> TCMT IO [Set Int] -> TCMT IO [Set Int]
forall a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> TCMT IO (Set Int)) -> [TAlt] -> TCMT IO [Set Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TAlt -> TCMT IO (Set Int)
goAlt [TAlt]
bs)
      TUnit{}   -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TSort{}   -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TErased{} -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TError{}  -> Set Int -> TCMT IO (Set Int)
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TCoerce TTerm
t -> TTerm -> TCMT IO (Set Int)
go TTerm
t

    goAlt :: TAlt -> TCMT IO (Set Int)
goAlt (TALit Literal
_   TTerm
b) = TTerm -> TCMT IO (Set Int)
go TTerm
b
    goAlt (TAGuard TTerm
g TTerm
b) = Set Int -> Set Int -> Set Int
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Set Int -> Set Int -> Set Int)
-> TCMT IO (Set Int) -> TCMT IO (Set Int -> Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCMT IO (Set Int)
go TTerm
g TCMT IO (Set Int -> Set Int)
-> TCMT IO (Set Int) -> TCMT IO (Set Int)
forall a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> TCMT IO (Set Int)
go TTerm
b
    goAlt (TACon QName
_ Int
a TTerm
b) = Int -> Set Int -> Set Int
forall {a}. (Num a, Ord a) => a -> Set a -> Set a
underBinders Int
a (Set Int -> Set Int) -> TCMT IO (Set Int) -> TCMT IO (Set Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCMT IO (Set Int)
go TTerm
b

    underBinder :: Set Int -> Set Int
underBinder = Int -> Set Int -> Set Int
forall {a}. (Num a, Ord a) => a -> Set a -> Set a
underBinders Int
1
    underBinders :: a -> Set a -> Set a
underBinders a
0 = Set a -> Set a
forall a. a -> a
id
    underBinders a
n = (a -> Bool) -> Set a -> Set a
forall a. (a -> Bool) -> Set a -> Set a
Set.filter (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
0) (Set a -> Set a) -> (Set a -> Set a) -> Set a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> Set a -> Set a
forall a b. (a -> b) -> Set a -> Set b
Set.mapMonotonic (a -> a -> a
forall a. Num a => a -> a -> a
subtract a
n)

stripUnusedArguments :: [ArgUsage] -> TTerm -> TTerm
stripUnusedArguments :: [ArgUsage] -> TTerm -> TTerm
stripUnusedArguments [ArgUsage]
used TTerm
t = Int -> TTerm -> TTerm
mkTLam Int
m (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Substitution' (SubstArg TTerm) -> TTerm -> TTerm
forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' TTerm
Substitution' (SubstArg TTerm)
rho TTerm
b
  where
    (Int
n, TTerm
b) = TTerm -> (Int, TTerm)
tLamView TTerm
t
    m :: Int
m      = [ArgUsage] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ArgUsage] -> Int) -> [ArgUsage] -> Int
forall a b. (a -> b) -> a -> b
$ (ArgUsage -> Bool) -> [ArgUsage] -> [ArgUsage]
forall a. (a -> Bool) -> [a] -> [a]
filter (ArgUsage -> ArgUsage -> Bool
forall a. Eq a => a -> a -> Bool
== ArgUsage
ArgUsed) [ArgUsage]
used'
    used' :: [ArgUsage]
used'  = [ArgUsage] -> [ArgUsage]
forall a. [a] -> [a]
reverse ([ArgUsage] -> [ArgUsage]) -> [ArgUsage] -> [ArgUsage]
forall a b. (a -> b) -> a -> b
$ Int -> [ArgUsage] -> [ArgUsage]
forall a. Int -> [a] -> [a]
take Int
n ([ArgUsage] -> [ArgUsage]) -> [ArgUsage] -> [ArgUsage]
forall a b. (a -> b) -> a -> b
$ [ArgUsage]
used [ArgUsage] -> [ArgUsage] -> [ArgUsage]
forall a. [a] -> [a] -> [a]
++ ArgUsage -> [ArgUsage]
forall a. a -> [a]
repeat ArgUsage
ArgUsed
    rho :: Substitution' TTerm
rho = [ArgUsage] -> Substitution' TTerm
computeSubst [ArgUsage]
used'
    computeSubst :: [ArgUsage] -> Substitution' TTerm
computeSubst (ArgUsage
ArgUnused : [ArgUsage]
bs) = TTerm
TErased TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a. a -> Substitution' a -> Substitution' a
:# [ArgUsage] -> Substitution' TTerm
computeSubst [ArgUsage]
bs
    computeSubst (ArgUsage
ArgUsed   : [ArgUsage]
bs) = Int -> Substitution' TTerm -> Substitution' TTerm
forall a. Int -> Substitution' a -> Substitution' a
liftS Int
1 (Substitution' TTerm -> Substitution' TTerm)
-> Substitution' TTerm -> Substitution' TTerm
forall a b. (a -> b) -> a -> b
$ [ArgUsage] -> Substitution' TTerm
computeSubst [ArgUsage]
bs
    computeSubst []               = Substitution' TTerm
forall a. Substitution' a
idS