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

import qualified Data.Set as Set

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

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

import Agda.Utils.Pretty (prettyShow)

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

computeUnused :: QName -> TTerm -> [Bool] -> TCM [Bool]
computeUnused :: QName -> TTerm -> [Bool] -> TCM [Bool]
computeUnused QName
q TTerm
t [Bool]
used = do
  VerboseKey -> Int -> VerboseKey -> TCMT IO ()
forall (m :: * -> *).
MonadDebug m =>
VerboseKey -> Int -> VerboseKey -> m ()
reportSLn VerboseKey
"treeless.opt.unused" Int
50 (VerboseKey -> TCMT IO ()) -> VerboseKey -> TCMT IO ()
forall a b. (a -> b) -> a -> b
$ VerboseKey
"Unused approximation for " VerboseKey -> VerboseKey -> VerboseKey
forall a. [a] -> [a] -> [a]
++ QName -> VerboseKey
forall a. Pretty a => a -> VerboseKey
prettyShow QName
q VerboseKey -> VerboseKey -> VerboseKey
forall a. [a] -> [a] -> [a]
++ VerboseKey
": " VerboseKey -> VerboseKey -> VerboseKey
forall a. [a] -> [a] -> [a]
++
                                       [VerboseKey] -> VerboseKey
unwords [ if Bool
u then [Char
x] else VerboseKey
"_" | (Char
x, Bool
u) <- VerboseKey -> [Bool] -> [(Char, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Char
'a'..] [Bool]
used ]
  QName -> [Bool] -> TCMT IO ()
setCompiledArgUse QName
q [Bool]
used
  Set Int
fv <- TTerm -> TCMT IO (Set Int)
go TTerm
t
  let used' :: [Bool]
used' = [ Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
i Set Int
fv | (Int
i, Bool
_) <- [(Int, Bool)] -> [(Int, Bool)]
forall a. [a] -> [a]
reverse ([(Int, Bool)] -> [(Int, Bool)]) -> [(Int, Bool)] -> [(Int, Bool)]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Bool] -> [(Int, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [Bool]
used ]
  if [Bool]
used [Bool] -> [Bool] -> Bool
forall a. Eq a => a -> a -> Bool
== [Bool]
used' then [Bool] -> TCM [Bool]
forall (m :: * -> *) a. Monad m => a -> m a
return [Bool]
used'
                   else QName -> TTerm -> [Bool] -> TCM [Bool]
computeUnused QName
q TTerm
t [Bool]
used'
  where
    go :: TTerm -> TCMT IO (Set Int)
go TTerm
t = case TTerm
t of
      TVar Int
x    -> Set Int -> TCMT IO (Set Int)
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 (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TDef{}    -> Set Int -> TCMT IO (Set Int)
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 (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TCon{}    -> Set Int -> TCMT IO (Set Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty

      TApp (TDef QName
f) Args
ts -> do
        [Bool]
used <- QName -> TCM [Bool]
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)
sequence [ TTerm -> TCMT IO (Set Int)
go TTerm
t | (TTerm
t, Bool
True) <- Args -> [Bool] -> [(TTerm, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip Args
ts ([Bool] -> [(TTerm, Bool)]) -> [Bool] -> [(TTerm, Bool)]
forall a b. (a -> b) -> a -> b
$ [Bool]
used [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
True ]

      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)
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 (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 (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)
mapM TAlt -> TCMT IO (Set Int)
goAlt [TAlt]
bs)
      TUnit{}   -> Set Int -> TCMT IO (Set Int)
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 (f :: * -> *) a. Applicative f => a -> f a
pure Set Int
forall a. Set a
Set.empty
      TErased{} -> Set Int -> TCMT IO (Set Int)
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 (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 (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 :: [Bool] -> TTerm -> TTerm
stripUnusedArguments :: [Bool] -> TTerm -> TTerm
stripUnusedArguments [Bool]
used TTerm
t = Int -> TTerm -> TTerm
mkTLam Int
m (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho TTerm
b
  where
    (Int
n, TTerm
b) = TTerm -> (Int, TTerm)
tLamView TTerm
t
    m :: Int
m      = [Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Bool] -> Int) -> [Bool] -> Int
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
forall a. a -> a
id [Bool]
used'
    used' :: [Bool]
used'  = [Bool] -> [Bool]
forall a. [a] -> [a]
reverse ([Bool] -> [Bool]) -> [Bool] -> [Bool]
forall a b. (a -> b) -> a -> b
$ Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take Int
n ([Bool] -> [Bool]) -> [Bool] -> [Bool]
forall a b. (a -> b) -> a -> b
$ [Bool]
used [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
True
    rho :: Substitution' TTerm
rho = [Bool] -> Substitution' TTerm
computeSubst [Bool]
used'
    computeSubst :: [Bool] -> Substitution' TTerm
computeSubst (Bool
False : [Bool]
bs) = TTerm
TErased TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a. a -> Substitution' a -> Substitution' a
:# [Bool] -> Substitution' TTerm
computeSubst [Bool]
bs
    computeSubst (Bool
True  : [Bool]
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
$ [Bool] -> Substitution' TTerm
computeSubst [Bool]
bs
    computeSubst []           = Substitution' TTerm
forall a. Substitution' a
idS