{-# OPTIONS_GHC -Wunused-imports #-}
module Mikan.Syntax.Scope.Trimming where

import Control.Monad.IO.Class
import Control.Monad

import Data.HashMap.Strict qualified as HMap
import Data.Map.Strict qualified as Map
import Data.Foldable
import Data.HashSet qualified as HSet
import Data.Set qualified as Set

import Mikan.Syntax.Common.Pretty
import Mikan.Syntax.Scope.State
import Mikan.Syntax.Scope.Base as A
import Mikan.Syntax.Abstract as A

import {-# SOURCE #-} Mikan.TypeChecking.Monad.Signature
import Mikan.TypeChecking.Monad.Statistics
import Mikan.TypeChecking.Monad.Debug
import Mikan.TypeChecking.Monad.State
import Mikan.TypeChecking.Monad.Base

import Mikan.Interaction.Options.ProfileOptions qualified as Profile

import Mikan.Utils.List1 qualified as List1
import Mikan.Utils.List1 (List1)
import Mikan.Utils.IORef.Strict
import Mikan.Utils.List

---------------------------------------------------------------------------
-- * Definition liveness & copy trimming
---------------------------------------------------------------------------

-- $liveness
-- The purpose of these operations is to let the type checker know what
-- definition copies are actually used for type-checking and need to be
-- created, and which can be safely skipped ("trimmed", hence the name
-- 'renTrimming').
--
-- The liveness the scope-checker reports should be a conservative
-- over-approximation of what names are actually used.  For example,
-- when an overloaded field/constructor name is resolved, we (have to)
-- mark all of the names in the set as used: not only do we not know
-- which one the type checker will actually pick, but the type checker
-- will ask for all of their
-- `Agda.TypeChecking.Monad.Signature.getConstInfo` when disambiguating.
--
-- Marking a name as live is done through the `markLiveName` function,
-- which is called from `Mikan.Syntax.Scope.Monad.resolveName'`, so any
-- names coming from concrete syntax will already be marked live.
-- However, if potential references to copies are being inserted
-- "synthetically" by the scope checker, they __must__ be referred to by
-- an explicit `markLiveName` call.
--
-- Module names can also be `markLiveName`d, and this has the effect of
-- transitively saving every definition having that module as a prefix.

-- | Create a new reference to store liveness information for names in
-- the given copied module.
newScopeCopyRef :: A.ModuleName -> ScopeM A.ScopeCopyRef
newScopeCopyRef :: ModuleName -> ScopeM ScopeCopyRef
newScopeCopyRef ModuleName
mname = IO ScopeCopyRef -> ScopeM ScopeCopyRef
forall a. IO a -> TCMT IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ScopeCopyRef -> ScopeM ScopeCopyRef)
-> IO ScopeCopyRef -> ScopeM ScopeCopyRef
forall a b. (a -> b) -> a -> b
$ ModuleName -> IORef LiveNames -> ScopeCopyRef
A.ScopeCopyRef ModuleName
mname (IORef LiveNames -> ScopeCopyRef)
-> IO (IORef LiveNames) -> IO ScopeCopyRef
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> LiveNames -> IO (IORef LiveNames)
forall a. a -> IO (IORef a)
newIORef LiveNames
forall a. Monoid a => a
mempty

-- | Mark a name as being a copy in the TC state, associating it with
-- the given 'A.ScopeCopyRef' for liveness information.
copyName :: A.ScopeCopyRef -> A.QName -> A.QName -> ScopeM ()
copyName :: ScopeCopyRef -> QName -> QName -> ScopeM ()
copyName !ScopeCopyRef
copy QName
from QName
to = do
  !from <- QName
-> ((ScopeCopyRef, QName) -> QName)
-> Maybe (ScopeCopyRef, QName)
-> QName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe QName
from (ScopeCopyRef, QName) -> QName
forall a b. (a, b) -> b
snd (Maybe (ScopeCopyRef, QName) -> QName)
-> (HashMap QName (ScopeCopyRef, QName)
    -> Maybe (ScopeCopyRef, QName))
-> HashMap QName (ScopeCopyRef, QName)
-> QName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QName
-> HashMap QName (ScopeCopyRef, QName)
-> Maybe (ScopeCopyRef, QName)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HMap.lookup QName
from (HashMap QName (ScopeCopyRef, QName) -> QName)
-> TCMT IO (HashMap QName (ScopeCopyRef, QName)) -> TCMT IO QName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Getter TCState (HashMap QName (ScopeCopyRef, QName))
-> TCMT IO (HashMap QName (ScopeCopyRef, QName))
forall (m :: * -> *) a. ReadTCState m => Getter TCState a -> m a
useTC (HashMap QName (ScopeCopyRef, QName)
 -> f (HashMap QName (ScopeCopyRef, QName)))
-> TCState -> f TCState
Lens' TCState (HashMap QName (ScopeCopyRef, QName))
Getter TCState (HashMap QName (ScopeCopyRef, QName))
stCopiedNames
  modifyingTC stCopiedNames $ HMap.insert to (copy, from)
  let
    k Maybe (HashSet QName)
Nothing  = HashSet QName -> Maybe (HashSet QName)
forall a. a -> Maybe a
Just (QName -> HashSet QName
forall a. Hashable a => a -> HashSet a
HSet.singleton QName
to)
    k (Just HashSet QName
s) = HashSet QName -> Maybe (HashSet QName)
forall a. a -> Maybe a
Just (QName -> HashSet QName -> HashSet QName
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HSet.insert QName
to HashSet QName
s)
  modifyingTC stNameCopies $ HMap.alter k from

-- | Class for entities which contain names that can be marked live.
-- The two fundamental instances are 'A.QName's and 'A.ModuleName's, but
-- there are some convenience instances for things like
-- 'A.ResolvedName's.
class MarkLive n where
  -- | Mark a name as (potentially) having been used.
  markLiveName :: n -> ScopeM ()

  default markLiveName :: forall t n'. (Traversable t, n ~ t n', MarkLive n') => n -> ScopeM ()
  markLiveName = (n' -> ScopeM ()) -> t n' -> ScopeM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ n' -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName

instance MarkLive n => MarkLive (List1 n)

instance MarkLive A.QName where
  markLiveName :: QName -> ScopeM ()
markLiveName QName
qn = QName
-> HashMap QName (ScopeCopyRef, QName)
-> Maybe (ScopeCopyRef, QName)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HMap.lookup QName
qn (HashMap QName (ScopeCopyRef, QName)
 -> Maybe (ScopeCopyRef, QName))
-> TCMT IO (HashMap QName (ScopeCopyRef, QName))
-> TCMT IO (Maybe (ScopeCopyRef, QName))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Getter TCState (HashMap QName (ScopeCopyRef, QName))
-> TCMT IO (HashMap QName (ScopeCopyRef, QName))
forall (m :: * -> *) a. ReadTCState m => Getter TCState a -> m a
useTC (HashMap QName (ScopeCopyRef, QName)
 -> f (HashMap QName (ScopeCopyRef, QName)))
-> TCState -> f TCState
Lens' TCState (HashMap QName (ScopeCopyRef, QName))
Getter TCState (HashMap QName (ScopeCopyRef, QName))
stCopiedNames TCMT IO (Maybe (ScopeCopyRef, QName))
-> (Maybe (ScopeCopyRef, QName) -> ScopeM ()) -> ScopeM ()
forall a b. TCMT IO a -> (a -> TCMT IO b) -> TCMT IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (A.ScopeCopyRef ModuleName
_ IORef LiveNames
ref, QName
_) ->
      IO () -> ScopeM ()
forall a. IO a -> TCMT IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ScopeM ()) -> IO () -> ScopeM ()
forall a b. (a -> b) -> a -> b
$ IORef LiveNames -> (LiveNames -> LiveNames) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef LiveNames
ref \case
        A.SomeLiveNames Set ModuleName
a Set QName
b -> Set ModuleName -> Set QName -> LiveNames
A.SomeLiveNames Set ModuleName
a (QName -> Set QName -> Set QName
forall a. Ord a => a -> Set a -> Set a
Set.insert QName
qn Set QName
b)
        LiveNames
A.AllLiveNames      -> LiveNames
A.AllLiveNames
    Maybe (ScopeCopyRef, QName)
Nothing -> () -> ScopeM ()
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

instance MarkLive A.ModuleName where
  markLiveName :: ModuleName -> ScopeM ()
markLiveName ModuleName
qn = Scope -> Maybe ScopeCopyRef
scopeIsCopy (Scope -> Maybe ScopeCopyRef)
-> TCMT IO Scope -> TCMT IO (Maybe ScopeCopyRef)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ModuleName -> TCMT IO Scope
getNamedScope ModuleName
qn TCMT IO (Maybe ScopeCopyRef)
-> (Maybe ScopeCopyRef -> ScopeM ()) -> ScopeM ()
forall a b. TCMT IO a -> (a -> TCMT IO b) -> TCMT IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just (A.ScopeCopyRef ModuleName
_ IORef LiveNames
ref) -> IO () -> ScopeM ()
forall a. IO a -> TCMT IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ScopeM ()) -> IO () -> ScopeM ()
forall a b. (a -> b) -> a -> b
$ IORef LiveNames -> (LiveNames -> LiveNames) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef LiveNames
ref \case
      A.SomeLiveNames Set ModuleName
a Set QName
b -> Set ModuleName -> Set QName -> LiveNames
A.SomeLiveNames (ModuleName -> Set ModuleName -> Set ModuleName
forall a. Ord a => a -> Set a -> Set a
Set.insert ModuleName
qn Set ModuleName
a) Set QName
b
      LiveNames
A.AllLiveNames      -> LiveNames
A.AllLiveNames
    Maybe ScopeCopyRef
Nothing -> () -> ScopeM ()
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

instance MarkLive A.AbstractName where
  markLiveName :: AbstractName -> ScopeM ()
markLiveName = QName -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName (QName -> ScopeM ())
-> (AbstractName -> QName) -> AbstractName -> ScopeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AbstractName -> QName
anameName

instance MarkLive ResolvedName where
  markLiveName :: ResolvedName -> ScopeM ()
markLiveName = \case
    ResolvedName
UnknownName         -> () -> ScopeM ()
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    VarName{}           -> () -> ScopeM ()
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    DefinedName Access
_ AbstractName
d Suffix
_   -> AbstractName -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName AbstractName
d
    FieldName List1 AbstractName
d         -> List1 AbstractName -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName List1 AbstractName
d
    ConstructorName Set1 Induction
_ List1 AbstractName
d -> List1 AbstractName -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName List1 AbstractName
d
    PatternSynResName List1 AbstractName
d -> List1 AbstractName -> ScopeM ()
forall n. MarkLive n => n -> ScopeM ()
markLiveName List1 AbstractName
d

-- | Read the 'LiveNames' from the given 'ScopeCopyRef'.
readLiveNames :: ScopeCopyRef -> ScopeM LiveNames
readLiveNames :: ScopeCopyRef -> ScopeM LiveNames
readLiveNames (ScopeCopyRef ModuleName
_ IORef LiveNames
ref) = IO LiveNames -> ScopeM LiveNames
forall a. IO a -> TCMT IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO LiveNames -> ScopeM LiveNames)
-> IO LiveNames -> ScopeM LiveNames
forall a b. (a -> b) -> a -> b
$ IORef LiveNames -> IO LiveNames
forall a. IORef a -> IO a
readIORef IORef LiveNames
ref

-- | Compare the trimming from an old 'ScopeCopyInfo' with a new one.
sameTrimming :: ScopeCopyInfo -> ScopeCopyInfo -> ScopeM Bool
sameTrimming :: ScopeCopyInfo -> ScopeCopyInfo -> ScopeM Bool
sameTrimming ScopeCopyInfo
old ScopeCopyInfo
new | ScopeCopyInfo -> Bool
renPublic ScopeCopyInfo
old, ScopeCopyInfo -> Bool
renPublic ScopeCopyInfo
new = Bool -> ScopeM Bool
forall a. a -> TCMT IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
sameTrimming ScopeCopyInfo
old ScopeCopyInfo
new = do
  a <- ScopeCopyRef -> ScopeM LiveNames
readLiveNames (ScopeCopyInfo -> ScopeCopyRef
renTrimming ScopeCopyInfo
old)
  b <- readLiveNames (renTrimming new)
  reportS "cache.trimming" 100 $ vcat
    [ "old trimming:" <+> pretty a
    , "new trimming:" <+> pretty b
    ]
  pure (a == b)

-- | Replaces the 'LiveNames' in every module application in the
-- scope-checker state at this point with 'AllLiveNames', effectively
-- disabling copy trimming for the modules we have seen so far.
clobberLiveNames :: ScopeM ()
clobberLiveNames :: ScopeM ()
clobberLiveNames = (Scope -> ScopeM ()) -> Map ModuleName Scope -> ScopeM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ((ScopeCopyRef -> ScopeM ()) -> Maybe ScopeCopyRef -> ScopeM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ScopeCopyRef -> ScopeM ()
forall {m :: * -> *}. MonadIO m => ScopeCopyRef -> m ()
go (Maybe ScopeCopyRef -> ScopeM ())
-> (Scope -> Maybe ScopeCopyRef) -> Scope -> ScopeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope -> Maybe ScopeCopyRef
scopeIsCopy) (Map ModuleName Scope -> ScopeM ())
-> (ScopeInfo -> Map ModuleName Scope) -> ScopeInfo -> ScopeM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScopeInfo -> Map ModuleName Scope
_scopeModules (ScopeInfo -> ScopeM ()) -> TCMT IO ScopeInfo -> ScopeM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TCMT IO ScopeInfo
forall (m :: * -> *). ReadTCState m => m ScopeInfo
getScope where
  go :: ScopeCopyRef -> m ()
go (ScopeCopyRef ModuleName
_ IORef LiveNames
ref) = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef LiveNames -> LiveNames -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef LiveNames
ref LiveNames
AllLiveNames

-- | Filter a 'ScopeCopyInfo' to only those names which were explicitly
-- referred to by the programmer.
--
-- Returns the new 'ScopeCopyInfo' and the complement of the definition
-- renaming, i.e. the names that were pruned.
onlyLiveCopies :: ModuleName -> ScopeCopyInfo -> TCM (Ren QName, ScopeCopyInfo)
onlyLiveCopies :: ModuleName -> ScopeCopyInfo -> TCM (Ren QName, ScopeCopyInfo)
onlyLiveCopies ModuleName
mn info :: ScopeCopyInfo
info@ScopeCopyInfo { renPublic :: ScopeCopyInfo -> Bool
renPublic = Bool
True } = (Ren QName
forall a. Monoid a => a
mempty, ScopeCopyInfo
info) (Ren QName, ScopeCopyInfo)
-> ScopeM () -> TCM (Ren QName, ScopeCopyInfo)
forall a b. a -> TCMT IO b -> TCMT IO a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$
  ProfileOption -> ScopeM () -> ScopeM ()
forall (m :: * -> *). MonadDebug m => ProfileOption -> m () -> m ()
whenProfile ProfileOption
Profile.Sections do
    VerboseKey -> ScopeM ()
forall (m :: * -> *). MonadStatistics m => VerboseKey -> m ()
tick    VerboseKey
"trimming: public copy"
    VerboseKey -> Word64 -> ScopeM ()
forall (m :: * -> *).
MonadStatistics m =>
VerboseKey -> Word64 -> m ()
tickMax VerboseKey
"largest public copy" (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Ren QName -> Int
forall a. Ren a -> Int
renamingSize (ScopeCopyInfo -> Ren QName
renNames ScopeCopyInfo
info)))

onlyLiveCopies ModuleName
mn info :: ScopeCopyInfo
info@ScopeCopyInfo { renNames :: ScopeCopyInfo -> Ren QName
renNames = Ren QName
rd, renTrimming :: ScopeCopyInfo -> ScopeCopyRef
renTrimming = ScopeCopyRef
ref } = do
  live <- ScopeCopyRef -> ScopeM LiveNames
readLiveNames ScopeCopyRef
ref

  reportS "tc.mod.apply.trim" 30 $ vcat
    [ "trimming renaming of module" <+> pretty mn
    , "  mods  =" <+> (case live of SomeLiveNames Set ModuleName
x Set QName
_ -> Set ModuleName -> Doc
forall a. Pretty a => a -> Doc
pretty Set ModuleName
x ; LiveNames
_ -> Doc
"*")
    , "  names =" <+> (case live of SomeLiveNames Set ModuleName
_ Set QName
x -> Set QName -> Doc
forall a. Pretty a => a -> Doc
pretty Set QName
x ; LiveNames
_ -> Doc
"*")
    , nest 2 (pretty info)
    ]

  if
    -- Trimming information from references to submodules (or
    -- module-level references to copy itself) is not propagated
    -- upwards, so if the copy is itself alive, everything will end up
    -- being copied --- we might as well skip the work of traversing the
    -- renaming.
    | mn `isModuleAlive` live -> (mempty, info) <$ whenProfile Profile.Sections do
      tick "trimming: live copy"

    | otherwise -> do
      let
      -- The scope checker will have marked any name from 'new' which is
      -- referred to as live (including, conservatively, every name from
      -- the submodules of 'new'), but the scope checker can't foresee
      -- which instances will be used --- and neither can we, here --- so
      -- we have to copy all of them, too.
        keep (QName
from, NonEmpty QName
to) = Definition -> Maybe InstanceInfo
defInstance (Definition -> Maybe InstanceInfo)
-> TCMT IO Definition -> TCMT IO (Maybe InstanceInfo)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> TCMT IO Definition
forall (m :: * -> *).
(HasConstInfo m, HasCallStack) =>
QName -> m Definition
getConstInfo QName
from TCMT IO (Maybe InstanceInfo)
-> (Maybe InstanceInfo
    -> Either (QName, NonEmpty QName) (QName, NonEmpty QName))
-> TCMT IO (Either (QName, NonEmpty QName) (QName, NonEmpty QName))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
          Just{} -> (QName, NonEmpty QName)
-> Either (QName, NonEmpty QName) (QName, NonEmpty QName)
forall a b. b -> Either a b
Right (QName
from, NonEmpty QName
to)
          Maybe InstanceInfo
_      -> case (QName -> Bool) -> NonEmpty QName -> [QName]
forall a. (a -> Bool) -> NonEmpty a -> [a]
List1.filter (QName -> LiveNames -> Bool
`isNameAlive` LiveNames
live) NonEmpty QName
to of
            []     -> (QName, NonEmpty QName)
-> Either (QName, NonEmpty QName) (QName, NonEmpty QName)
forall a b. a -> Either a b
Left (QName
from, NonEmpty QName
to)
            (QName
x:[QName]
xs) -> (QName, NonEmpty QName)
-> Either (QName, NonEmpty QName) (QName, NonEmpty QName)
forall a b. b -> Either a b
Right (QName
from, QName
x QName -> [QName] -> NonEmpty QName
forall a. a -> [a] -> NonEmpty a
List1.:| [QName]
xs)

        rsz = Ren QName -> Int
forall a. Ren a -> Int
renamingSize Ren QName
rd

      (deleted, kept) <- partitionEithers' <$> traverse keep (Map.toAscList rd)
      let
        !rd   = [(QName, NonEmpty QName)] -> Ren QName
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList [(QName, NonEmpty QName)]
kept
        saved = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
rsz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Ren QName -> Int
forall a. Ren a -> Int
renamingSize Ren QName
rd)

      whenProfile Profile.Sections $ tickN "trimmed definitions" saved
      when (saved == 0) do
        reportSLn "tc.mod.apply.trim" 30 "... but nothing happened!"
        whenProfile Profile.Sections $ tick "trimming: no effect"

      let
        -- Elaboration of overloaded projections depends on whether the
        -- section is in the signature or not, and since copying
        -- sections is cheap (copying definitions is the big issue),
        -- it's easier to just preserve all the original sections than
        -- it is to trim them.
        info' = ScopeCopyInfo
info { renNames = rd }
      (Map.fromAscList deleted, info') <$ reportS "tc.mod.apply.trim" 30 (pretty info')