Erase.hs

{-# OPTIONS_GHC -Wunused-imports #-}

{-# LANGUAGE PatternSynonyms #-}

module Agda.Compiler.Treeless.Erase
       ( eraseTerms
       , computeErasedConstructorArgs
       , isErasable
       ) where

import Control.Arrow       ( first, second )
import Control.Monad.State ( StateT, evalStateT )

import Data.Map (Map)
import qualified Data.Map as Map

import Agda.Syntax.Common
import Agda.Syntax.Internal as I
import Agda.Syntax.Treeless
import Agda.Syntax.Literal

import Agda.TypeChecking.Substitute
import Agda.TypeChecking.Monad as I
import Agda.TypeChecking.Telescope
import Agda.TypeChecking.Datatypes
import Agda.TypeChecking.Pretty
import Agda.TypeChecking.Primitive

import Agda.Compiler.Treeless.Subst
import Agda.Compiler.Treeless.Unused

import Agda.Utils.Functor
import Agda.Utils.Lens
import Agda.Utils.Maybe
import Agda.Utils.Memo
import Agda.Utils.Monad
import Agda.Syntax.Common.Pretty (prettyShow)
import Agda.Utils.IntSet.Infinite (IntSet)
import qualified Agda.Utils.IntSet.Infinite as IntSet

import Agda.Utils.Impossible

-- | State of the eraser.
data ESt = ESt
  { _funMap  :: Map QName FunInfo
      -- ^ Memoize computed `FunInfo` for functions/constructors/... `QName`.
  , _typeMap :: Map QName TypeInfo
      -- ^ Memoize computed `TypeInfo` for data/record types `QName`.
  }

funMap :: Lens' ESt (Map QName FunInfo)
funMap f r = f (_funMap r) <&> \ a -> r { _funMap = a }

typeMap :: Lens' ESt (Map QName TypeInfo)
typeMap f r = f (_typeMap r) <&> \ a -> r { _typeMap = a }

-- | Eraser monad.
type E = StateT ESt TCM

runE :: E a -> TCM a
runE m = evalStateT m (ESt Map.empty Map.empty)

-- | Takes the name of the data/record type.
computeErasedConstructorArgs :: QName -> TCM ()
computeErasedConstructorArgs d = do
  cs <- getNotErasedConstructors d
  runE $ mapM_ getFunInfo cs

eraseTerms :: QName -> EvaluationStrategy -> TTerm -> TCM TTerm
eraseTerms q eval t = usedArguments q t *> runE (eraseTop q t)
  where
    eraseTop q t = do
      (_, h) <- getFunInfo q
      case h of
        Erasable -> pure TErased
        Empty    -> pure TErased
        _        -> erase t

    erase t = case tAppView t of

      (TCon c, vs) -> do
        (rs, h) <- getFunInfo c
        when (length rs < length vs) __IMPOSSIBLE__
        case h of
          Erasable -> pure TErased
          Empty    -> pure TErased
          _        -> tApp (TCon c) <$> zipWithM eraseRel rs vs

      (TDef f, vs) -> do
        (rs, h) <- getFunInfo f
        case h of
          Erasable -> pure TErased
          Empty    -> pure TErased
          _        -> tApp (TDef f) <$> zipWithM eraseRel (rs ++ repeat NotErasable) vs

      _ -> case t of
        TVar{}         -> pure t
        TDef{}         -> pure t
        TPrim{}        -> pure t
        TLit{}         -> pure t
        TCon{}         -> pure t
        TApp f es      -> tApp <$> erase f <*> mapM erase es
        TLam b         -> tLam <$> erase b
        TLet e b       -> do
          e <- erase e
          if isErased e
            then case b of
                   TCase 0 _ _ _ -> tLet TErased <$> erase b
                   _             -> erase $ subst 0 TErased b
            else tLet e <$> erase b
        TCase x t d bs -> do
          (d, bs) <- pruneUnreachable x (caseErased t) (caseType t) d bs
          d       <- erase d
          bs      <- mapM eraseAlt bs
          tCase x t d bs

        TUnit          -> pure t
        TSort          -> pure t
        TErased        -> pure t
        TError{}       -> pure t
        TCoerce e      -> TCoerce <$> erase e

    -- #3380: this is not safe for strict backends
    tLam TErased | eval == LazyEvaluation = TErased
    tLam t                                = TLam t

    tLet e b
      | freeIn 0 b = TLet e b
      | otherwise  = strengthen impossible b

    tApp f []                  = f
    tApp TErased _             = TErased
    tApp f _ | isUnreachable f = tUnreachable
    tApp f es                  = mkTApp f es

    tCase x t d bs
      | isErased d && all (isErased . aBody) bs = pure TErased
      | otherwise = case bs of
        [b@(TACon c _ _)] -> do
          h <- snd <$> getFunInfo c
          case h of
            NotErasable -> fallback
            Empty       -> pure TErased
            Erasable    -> erasedBody b
        _ -> fallback
      where
        noerase = pure $ TCase x t d bs

        erasedBody = \case
          TACon _ arity body ->
            (if arity == 0 then pure else erase) $
               -- might enable more erasure
            applySubst (replicate arity TErased ++# idS) body
          TALit _ body   -> pure body
          TAGuard _ body -> pure body

        fallback = case (caseErased t, bs) of
          (Erased{}, [b]) ->
            -- The case variable is erased, and there is exactly one
            -- case: use the case's body.
            erasedBody b
          (Erased{}, [])  ->
            -- The case variable is erased, and there is no case: use
            -- the default.
            pure $ if isErased d then TErased else d
          (Erased{}, _ : _ : _) ->
            -- The case variable is erased, and there are at least two
            -- cases: crash.
            __IMPOSSIBLE__
          _ ->
            -- The case variable is not erased: do not erase anything.
            noerase

    isErased t = t == TErased || isUnreachable t

    eraseRel r t | erasable r = pure TErased
                 | otherwise  = erase t

    eraseAlt = \case
      TALit l b   -> TALit l   <$> erase b
      TACon c a b -> do
        rs <- map erasable . fst <$> getFunInfo c
        let sub = foldr (\ e -> if e then (TErased :#) . wkS 1 else liftS 1) idS $ reverse rs
        TACon c a <$> erase (applySubst sub b)
      TAGuard g b -> TAGuard   <$> erase g <*> erase b

pruneUnreachable ::
  Int -> Erased -> CaseType -> TTerm -> [TAlt] -> E (TTerm, [TAlt])
pruneUnreachable x erased t d bs = case erased of
  NotErased{} -> pruneUnreachable' x erased t d bs
  Erased{}    ->
    -- If the match is on an erased argument, then the first branch
    -- should match.
    case bs of
      []    -> pruneUnreachable' x erased t d            []
      b : _ -> pruneUnreachable' x erased t tUnreachable [b]

-- | Doesn't have any type information (other than the name of the data type),
--   so we can't do better than checking if all constructors are present.
pruneUnreachable' ::
  Int -> Erased -> CaseType -> TTerm -> [TAlt] -> E (TTerm, [TAlt])
pruneUnreachable' _ erased (CTData q) d bs' = do
  -- Erased constructors are pruned iff the match is made on a
  -- non-erased argument.
  cs <- lift $
        if isErased erased
        then getConstructors q
        else getNotErasedConstructors q
  let bs | isErased erased = bs'
         | otherwise       =
           flip filter bs' $ \case
             a@TACon{} -> (aCon a) `elem` cs
             TAGuard{} -> True
             TALit{}   -> True
  let -- In the case of a match on an erased argument the value d is
      -- equal to tUnreachable, except perhaps if bs is empty. In the
      -- latter case complete is True exactly when the type has zero
      -- constructors (erased or not), in which case it makes sense to
      -- replace d with tUnreachable.
      complete = length cs == length [ b | b@TACon{} <- bs ]
      d' | complete  = tUnreachable
         | otherwise = d
  return (d', bs)
pruneUnreachable' x _ CTNat d bs =
  return $ pruneIntCase x d bs (IntSet.below 0)
pruneUnreachable' x _ CTInt d bs =
  return $ pruneIntCase x d bs IntSet.empty
pruneUnreachable' _ _ _ d bs =
  pure (d, bs)

-- These are the guards we generate for Int/Nat pattern matching
pattern Below :: Int -> Integer -> TTerm
pattern Below x n = TApp (TPrim PLt)  [TVar x, TLit (LitNat n)]

pattern Above :: Int -> Integer -> TTerm
pattern Above x n = TApp (TPrim PGeq) [TVar x, TLit (LitNat n)]

-- | Strip unreachable clauses (replace by tUnreachable for the default).
--   Fourth argument is the set of ints covered so far.
pruneIntCase :: Int -> TTerm -> [TAlt] -> IntSet -> (TTerm, [TAlt])
pruneIntCase x d bs cover = go bs cover
  where
    go [] cover
      | cover == IntSet.full = (tUnreachable, [])
      | otherwise            = (d, [])
    go (b : bs) cover =
      case b of
        TAGuard (Below y n) _ | x == y -> rec (IntSet.below n)
        TAGuard (Above y n) _ | x == y -> rec (IntSet.above n)
        TALit (LitNat n) _             -> rec (IntSet.singleton n)
        _                                -> second (b :) $ go bs cover
      where
        rec this = second addAlt $ go bs cover'
          where
            this'  = IntSet.difference this cover
            cover' = this' <> cover
            addAlt = case IntSet.toFiniteList this' of
                       Just []  -> id                                     -- unreachable case
                       Just [n] -> (TALit (LitNat n) (aBody b) :) -- possibly refined case
                       _        -> (b :)                                  -- unchanged case

data TypeInfo = Empty | Erasable | NotErasable
  deriving (Eq, Show)

sumTypeInfo :: [TypeInfo] -> TypeInfo
sumTypeInfo is = foldr plus Empty is
  where
    plus Empty       r           = r
    plus r           Empty       = r
    plus Erasable    r           = r
    plus r           Erasable    = r
    plus NotErasable NotErasable = NotErasable

erasable :: TypeInfo -> Bool
erasable Erasable    = True
erasable Empty       = True
erasable NotErasable = False

type FunInfo = ([TypeInfo], TypeInfo)

getFunInfo :: QName -> E FunInfo
getFunInfo q = memo (funMap . key q) $ getInfo q
  where
    getInfo :: QName -> E FunInfo
    getInfo q = do
      (rs, t) <- do
        (tel, t) <- lift $ typeWithoutParams q
        is     <- mapM (getTypeInfo . snd . dget) tel
        used   <- lift $ (++ repeat ArgUsed) . fromMaybe [] <$> getCompiledArgUse q
        forced <- lift $ (++ repeat NotForced) <$> getForcedArgs q
        return (zipWith3 (uncurry . mkR . getModality) tel (zip forced used) is, t)
      h <- if isAbsurdLambdaName q then pure Erasable else getTypeInfo t
      lift $ reportSLn "treeless.opt.erase.info" 50 $ "type info for " ++ prettyShow q ++ ": " ++ show rs ++ " -> " ++ show h
      lift $ setErasedConArgs q $ map erasable rs
      return (rs, h)

    -- Treat empty, erasable, or unused arguments as Erasable
    mkR :: Modality -> IsForced -> ArgUsage -> TypeInfo -> TypeInfo
    mkR m f u i
      | not (usableModality m) = Erasable
      | ArgUnused <- u         = Erasable
      | Forced <- f            = Erasable
      | otherwise              = i

isErasable :: QName -> TCM Bool
isErasable qn =
  erasable . snd <$> runE (getFunInfo qn)

telListView :: Type -> TCM (ListTel, Type)
telListView t = do
  TelV tel t <- telViewPath t
  return (telToList tel, t)

typeWithoutParams :: QName -> TCM (ListTel, Type)
typeWithoutParams q = do
  def <- getConstInfo q
  let d = case I.theDef def of
        Function{ funProjection = Right Projection{ projIndex = i } } -> i - 1
        Constructor{ conPars = n } -> n
        _                          -> 0
  first (drop d) <$> telListView (defType def)

getTypeInfo :: Type -> E TypeInfo
getTypeInfo t0 = do
  (tel, t) <- lift $ telListView t0
  et <- case I.unEl t of
    I.Def d _ -> do
      -- #2916: Only update the memo table for d. Results for other types are
      -- under the assumption that d is erasable!
      oldMap <- use typeMap
      dInfo <- typeInfo d
      typeMap .= Map.insert d dInfo oldMap
      return dInfo
    Sort{}    -> return Erasable
    _         -> return NotErasable
  is <- mapM (getTypeInfo . snd . dget) tel
  let e | Empty `elem` is = Erasable
        | null is         = et        -- TODO: guard should really be "all inhabited is"
        | et == Empty     = Erasable
        | otherwise       = et
  lift $ reportSDoc "treeless.opt.erase.type" 50 $ prettyTCM t0 <+> text ("is " ++ show e)
  return e
  where
  typeInfo :: QName -> E TypeInfo
  typeInfo q = ifM (erasureForbidden q) (return NotErasable) $ {-else-} do
    memoRec (typeMap . key q) Erasable $ do  -- assume recursive occurrences are erasable
      msizes <- lift $ mapM getBuiltinName
                         [builtinSize, builtinSizeLt]
      def    <- lift $ getConstInfo q
      let mcs = case I.theDef def of
                  I.Datatype{ dataCons = cs } -> Just cs
                  I.Record{ recConHead = c }  -> Just [conName c]
                  _                           -> Nothing
      case mcs of
        _ | Just q `elem` msizes -> return Erasable
        Just [c] -> do
          (ts, _) <- lift $ typeWithoutParams c
          let rs = map getModality ts
          is <- mapM (getTypeInfo . snd . dget) ts
          let er = and [ erasable i || not (usableModality r) | (i, r) <- zip is rs ]
          return $ if er then Erasable else NotErasable
        Just []      -> return Empty
        Just (_:_:_) -> return NotErasable
        Nothing ->
          case I.theDef def of
            I.Function{ funClauses = cs } ->
              sumTypeInfo <$> mapM (maybe (return Empty) (getTypeInfo . El __DUMMY_SORT__) . clauseBody) cs
            _ -> return NotErasable
  -- The backend also has a say whether a type is eraseable or not.
  erasureForbidden :: QName -> E Bool
  erasureForbidden q = lift $ not <$> activeBackendMayEraseType q