{-# OPTIONS_GHC -Wunused-imports #-}

module Mikan.TypeChecking.DeadCode (eliminateDeadCode) where

import Control.Monad
import Control.Monad.Trans

import Data.Maybe
import Data.Map.Strict qualified as MapS
import Data.HashMap.Strict qualified as HMap
import Data.IORef

import Mikan.Syntax.Common
import Mikan.Syntax.Internal
import Mikan.Syntax.Internal.Names
import Mikan.Syntax.Scope.Base

import Mikan.Benchmarking qualified as Bench
import Mikan.TypeChecking.Monad.Benchmark qualified as Bench

import Mikan.TypeChecking.Monad

import Mikan.Utils.Impossible
import Mikan.Utils.List
import Mikan.Utils.List1 qualified as List1

import Mikan.Utils.HashTable (HashTableLU, HashTableLL)
import Mikan.Utils.HashTable qualified as HT

-- [TODO: Reed M, 30/04/2026] Investigate if its worth computing precise
-- reachability now that rewrite rules are gone.

-- | Run before serialisation to remove data that's not reachable from the
-- public interface.
--
-- We do not compute reachable data precisely, because that
-- would be very expensive. This was originally due to rewrite
-- rules, though this may no longer be true after support for @--rewriting@
-- was dropped in Mikan.
--
-- The following things are assumed to be "roots":
--   - public definitions
--   - definitions marked as primitive
--   - definitions with COMPILE pragma
--   - all pattern synonyms (because currently all of them go into interfaces)
--   - all parameter sections (because currently all of them go into interfaces)
--     (see also issues #6931 and #7382)
--   - local builtins
--   - closed display forms
-- We only ever prune dead metavariables and definitions. We return the pruned metas,
-- pruned definitions and closed display forms.
eliminateDeadCode :: ScopeInfo -> TCM (RemoteMetaStore, Definitions, DisplayForms)
eliminateDeadCode :: ScopeInfo
-> TCM
     (RemoteMetaStore, Definitions,
      HashMap QName (NonEmpty (Open DisplayForm)))
eliminateDeadCode !ScopeInfo
scope = Account (BenchPhase (TCMT IO))
-> TCM
     (RemoteMetaStore, Definitions,
      HashMap QName (NonEmpty (Open DisplayForm)))
-> TCM
     (RemoteMetaStore, Definitions,
      HashMap QName (NonEmpty (Open DisplayForm)))
forall (m :: * -> *) c.
MonadBench m =>
Account (BenchPhase m) -> m c -> m c
Bench.billTo [BenchPhase (TCMT IO)
Phase
Bench.DeadCode] (TCM
   (RemoteMetaStore, Definitions,
    HashMap QName (NonEmpty (Open DisplayForm)))
 -> TCM
      (RemoteMetaStore, Definitions,
       HashMap QName (NonEmpty (Open DisplayForm))))
-> TCM
     (RemoteMetaStore, Definitions,
      HashMap QName (NonEmpty (Open DisplayForm)))
-> TCM
     (RemoteMetaStore, Definitions,
      HashMap QName (NonEmpty (Open DisplayForm)))
forall a b. (a -> b) -> a -> b
$ do
  !sig <- TCMT IO Signature
forall (m :: * -> *). ReadTCState m => m Signature
getSignature
  let !defs = Signature
sig Signature
-> Getting Definitions Signature Definitions -> Definitions
forall s a. s -> Getting a s a -> a
^. Getting Definitions Signature Definitions
Lens' Signature Definitions
sigDefinitions
  !metas <- useR stSolvedMetaStore

  -- #2921: Eliminating definitions with attached COMPILE pragmas results in
  -- the pragmas not being checked. Simple solution: don't eliminate these.
  -- #6022 (Andreas, 2022-09-30): Eliminating cubical primitives can lead to crashes.
  -- Simple solution: retain all primitives (shouldn't be many).
  let hasCompilePragma = Bool -> Bool
not (Bool -> Bool) -> (Definition -> Bool) -> Definition -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map BackendName [CompilerPragma] -> Bool
forall k a. Map k a -> Bool
MapS.null (Map BackendName [CompilerPragma] -> Bool)
-> (Definition -> Map BackendName [CompilerPragma])
-> Definition
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Definition -> Map BackendName [CompilerPragma]
defCompiledRep

      isPrimitive = \case
        Primitive{}     -> Bool
True
        PrimitiveSort{} -> Bool
True
        Defn
_               -> Bool
False

      extraRootsFilter (a
name, Definition
def)
        | Definition -> Bool
hasCompilePragma Definition
def Bool -> Bool -> Bool
|| Defn -> Bool
isPrimitive (Definition -> Defn
theDef Definition
def) = a -> Maybe a
forall a. a -> Maybe a
Just a
name
        | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing

  let !pubModules = ScopeInfo -> Map ModuleName Scope
publicModules ScopeInfo
scope

    -- Ulf, 2016-04-12:
    -- Non-closed display forms are not applicable outside the module anyway,
    -- and should be dead-code eliminated (#1928).
  let eliminateNonClosed = [Open a] -> Maybe (NonEmpty (Open a))
forall a. [a] -> Maybe (NonEmpty a)
List1.nonEmpty ([Open a] -> Maybe (NonEmpty (Open a)))
-> (NonEmpty (Open a) -> [Open a])
-> NonEmpty (Open a)
-> Maybe (NonEmpty (Open a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Open a -> Bool) -> NonEmpty (Open a) -> [Open a]
forall a. (a -> Bool) -> NonEmpty a -> [a]
List1.filter Open a -> Bool
forall a. Open a -> Bool
isClosed

  !rootDisplayForms <-
      HMap.mapMaybe eliminateNonClosed <$> useTC stImportsDisplayForms

  let !rootPubNames  = (AbstractName -> QName) -> [AbstractName] -> [QName]
forall a b. (a -> b) -> [a] -> [b]
map' AbstractName -> QName
anameName ([AbstractName] -> [QName]) -> [AbstractName] -> [QName]
forall a b. (a -> b) -> a -> b
$ Map ModuleName Scope -> [AbstractName]
publicNamesOfModules Map ModuleName Scope
pubModules
  let !rootExtraDefs = ((QName, Definition) -> Maybe QName)
-> [(QName, Definition)] -> [QName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (QName, Definition) -> Maybe QName
forall {a}. (a, Definition) -> Maybe a
extraRootsFilter ([(QName, Definition)] -> [QName])
-> [(QName, Definition)] -> [QName]
forall a b. (a -> b) -> a -> b
$ Definitions -> [(QName, Definition)]
forall k v. HashMap k v -> [(k, v)]
HMap.toList Definitions
defs
  let !rootModSections = Signature
sig Signature -> Getting Sections Signature Sections -> Sections
forall s a. s -> Getting a s a -> a
^. Getting Sections Signature Sections
Lens' Signature Sections
sigSections
  !rootBuiltins <- useTC stLocalBuiltins
  !rootPatSyns  <- getPatternSyns

  -- When we're traversing the definition of a name, we want to record if the definition contains
  -- any metas. If a definition doesn't contain metas, we can skip instantiateFull on it before
  -- serialization.
  --
  -- We record "meta-closedness" in the following way. If we're traversing a name definition,
  -- "insideDef" contains the IORef Bool which can be flipped if we hit a metavariable. We
  -- store these IORef Bool-s in "seenNames". A new IORef Bool is created for each QName when we
  -- encounter it for the first time.
  !insideDef <- lift $ newIORef Nothing :: TCM (IORef (Maybe (IORef Bool)))

  !seenNames <- lift $ HT.empty :: TCM (HashTableLL QName (IORef Bool))
  !seenMetas <- lift $ HT.empty :: TCM (HashTableLU MetaId ())

  let goName :: QName -> IO ()
      goName !QName
x = HashTableLL QName (IORef Bool)
-> QName
-> (IORef Bool -> IO ())
-> IO (IORef Bool)
-> (IORef Bool -> IO ())
-> IO ()
forall (ks :: * -> * -> *) k (vs :: * -> * -> *) v a.
(Hashable k, MVector ks k, MVector vs v) =>
HashTable ks k vs v
-> k -> (v -> IO a) -> IO v -> (v -> IO a) -> IO a
HT.insertingIfAbsent HashTableLL QName (IORef Bool)
seenNames QName
x
        -- already visited name, do nothing
        (\IORef Bool
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

        -- unvisited, insert new IORef into table
        (Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False)

        (\IORef Bool
ref -> do
            prevRef <- IORef (Maybe (IORef Bool)) -> IO (Maybe (IORef Bool))
forall a. IORef a -> IO a
readIORef IORef (Maybe (IORef Bool))
insideDef  -- save ref
            writeIORef insideDef (Just ref)
            go (HMap.lookup x defs)
            writeIORef insideDef prevRef    -- restore ref
        )

      goMeta :: MetaId -> IO ()
      goMeta !MetaId
m = do
        prevRef <- IORef (Maybe (IORef Bool)) -> IO (Maybe (IORef Bool))
forall a. IORef a -> IO a
readIORef IORef (Maybe (IORef Bool))
insideDef
        case prevRef of
          -- we're not inside any name definition
          Maybe (IORef Bool)
Nothing  -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          -- we're inside a name definition, record meta occurrence
          Just IORef Bool
ref -> IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
ref Bool
True

        HT.insertingIfAbsent seenMetas m
          -- already visited, do nothing
          (\()
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

          (pure ())
          -- unvisited
          (\()
_ -> case MetaId -> Map MetaId MetaVariable -> Maybe MetaVariable
forall k a. Ord k => k -> Map k a -> Maybe a
MapS.lookup MetaId
m Map MetaId MetaVariable
metas of
            Maybe MetaVariable
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just MetaVariable
mv -> do
              IORef (Maybe (IORef Bool)) -> Maybe (IORef Bool) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IORef Bool))
insideDef Maybe (IORef Bool)
forall a. Maybe a
Nothing
              Term -> IO ()
forall a. NamesIn a => a -> IO ()
go (Instantiation -> Term
instBody (MetaVariable -> Instantiation
theInstantiation MetaVariable
mv))
              Teletype -> IO ()
forall a. NamesIn a => a -> IO ()
go (Judgement MetaId -> Teletype
forall a. Judgement a -> Teletype
jMetaType (MetaVariable -> Judgement MetaId
mvJudgement MetaVariable
mv))
              IORef (Maybe (IORef Bool)) -> Maybe (IORef Bool) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IORef Bool))
insideDef Maybe (IORef Bool)
prevRef -- restore ref
          )

      go :: NamesIn a => a -> IO ()
      go !a
x = (Either QName MetaId -> IO ()) -> a -> IO ()
forall m. Monoid m => (Either QName MetaId -> m) -> a -> m
forall a m.
(NamesIn a, Monoid m) =>
(Either QName MetaId -> m) -> a -> m
namesAndMetasIn' ((QName -> IO ())
-> (MetaId -> IO ()) -> Either QName MetaId -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either QName -> IO ()
goName MetaId -> IO ()
goMeta) a
x
      {-# INLINE go #-}

  Bench.billTo [Bench.DeadCode, Bench.DeadCodeReachable] $ lift $ do
    go rootDisplayForms
    foldMap goName rootPubNames
    foldMap goName rootExtraDefs
    go rootModSections
    go rootBuiltins
    foldMap go rootPatSyns

  let goMetas :: [(MetaId, MetaVariable)] -> IO [(MetaId, RemoteMetaVariable)]
      goMetas [] = [(MetaId, RemoteMetaVariable)] -> IO [(MetaId, RemoteMetaVariable)]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
      goMetas ((!MetaId
i, !MetaVariable
m):[(MetaId, MetaVariable)]
metas) = HashTableLU MetaId () -> MetaId -> IO (Maybe ())
forall k (ks :: * -> * -> *) (vs :: * -> * -> *) v.
(Hashable k, MVector ks k, MVector vs v) =>
HashTable ks k vs v -> k -> IO (Maybe v)
HT.lookup HashTableLU MetaId ()
seenMetas MetaId
i IO (Maybe ())
-> (Maybe () -> IO [(MetaId, RemoteMetaVariable)])
-> IO [(MetaId, RemoteMetaVariable)]
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe ()
Nothing -> [(MetaId, MetaVariable)] -> IO [(MetaId, RemoteMetaVariable)]
goMetas [(MetaId, MetaVariable)]
metas
        Just ()
_  -> do
          let !m' :: RemoteMetaVariable
m' = MetaVariable -> RemoteMetaVariable
remoteMetaVariable MetaVariable
m
          ((MetaId
i, RemoteMetaVariable
m')(MetaId, RemoteMetaVariable)
-> [(MetaId, RemoteMetaVariable)] -> [(MetaId, RemoteMetaVariable)]
forall a. a -> [a] -> [a]
:) ([(MetaId, RemoteMetaVariable)] -> [(MetaId, RemoteMetaVariable)])
-> IO [(MetaId, RemoteMetaVariable)]
-> IO [(MetaId, RemoteMetaVariable)]
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> [(MetaId, MetaVariable)] -> IO [(MetaId, RemoteMetaVariable)]
goMetas [(MetaId, MetaVariable)]
metas

      goDefs :: [(QName, Definition)] -> IO [(QName, Definition)]
      goDefs [] = [(QName, Definition)] -> IO [(QName, Definition)]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
      goDefs ((!QName
x, !Definition
d):[(QName, Definition)]
defs) = HashTableLL QName (IORef Bool) -> QName -> IO (Maybe (IORef Bool))
forall k (ks :: * -> * -> *) (vs :: * -> * -> *) v.
(Hashable k, MVector ks k, MVector vs v) =>
HashTable ks k vs v -> k -> IO (Maybe v)
HT.lookup HashTableLL QName (IORef Bool)
seenNames QName
x IO (Maybe (IORef Bool))
-> (Maybe (IORef Bool) -> IO [(QName, Definition)])
-> IO [(QName, Definition)]
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe (IORef Bool)
Nothing -> [(QName, Definition)] -> IO [(QName, Definition)]
goDefs [(QName, Definition)]
defs
        Just IORef Bool
r  -> do
          hasM <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
r
          !d <- case hasM of
            Bool
True  -> Definition -> IO Definition
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Definition
d
            Bool
False -> Definition -> IO Definition
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Definition -> IO Definition) -> Definition -> IO Definition
forall a b. (a -> b) -> a -> b
$! Definition
d {defMightContainMetas = False}
          ((x, d):) <$!> goDefs defs

  !metas <- lift $ HMap.fromList <$> goMetas (MapS.toList metas)
  !defs  <- lift $ HMap.fromList <$> goDefs (HMap.toList defs)
  pure (metas, defs, rootDisplayForms)

-- | Returns the instantiation.
--   Precondition: The instantiation must be of the form @'InstV' inst@.
theInstantiation :: MetaVariable -> Instantiation
theInstantiation :: MetaVariable -> Instantiation
theInstantiation MetaVariable
mv = case MetaVariable -> MetaInstantiation
mvInstantiation MetaVariable
mv of
  InstV Instantiation
inst                     -> Instantiation
inst
  OpenMeta{}                     -> Instantiation
forall a. HasCallStack => a
__IMPOSSIBLE__
  BlockedConst{}                 -> Instantiation
forall a. HasCallStack => a
__IMPOSSIBLE__
  PostponedTypeCheckingProblem{} -> Instantiation
forall a. HasCallStack => a
__IMPOSSIBLE__

-- | Converts from 'MetaVariable' to 'RemoteMetaVariable'.
--   Precondition: The instantiation must be of the form @'InstV' inst@.
remoteMetaVariable :: MetaVariable -> RemoteMetaVariable
remoteMetaVariable :: MetaVariable -> RemoteMetaVariable
remoteMetaVariable !MetaVariable
mv = RemoteMetaVariable
  { rmvInstantiation :: Instantiation
rmvInstantiation = MetaVariable -> Instantiation
theInstantiation MetaVariable
mv
  , rmvJudgement :: Judgement MetaId
rmvJudgement     = MetaVariable -> Judgement MetaId
mvJudgement MetaVariable
mv
  }