content: {-# LANGUAGE FlexibleInstances          #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE LambdaCase                 #-}
  {-# LANGUAGE MultiWayIf                 #-}
  {-# LANGUAGE OverloadedStrings          #-}
  
  -- | Two space stop-and-copy garbage collector.
  --
  -- * Compacting: memory is moved to a new location
  -- * Tracing
  module Stg.Machine.GarbageCollection.TwoSpaceCopying (
      twoSpaceCopying,
  ) where
  
  
  
  import           Control.Monad.Trans.Class
  import           Control.Monad.Trans.Reader
  import           Control.Monad.Trans.State
  import           Data.Foldable
  import           Data.Map                   (Map)
  import qualified Data.Map                   as M
  import           Data.Sequence              (Seq, ViewL (..), (|>))
  import qualified Data.Sequence              as Seq
  import           Data.Set                   (Set)
  import qualified Data.Set                   as S
  
  import           Stg.Machine.GarbageCollection.Common
  import qualified Stg.Machine.Heap                     as H
  import           Stg.Machine.Types
  
  
  
  -- | Move all used addresses by moving them to a safe location, and delete the
  -- leftovers.
  twoSpaceCopying :: GarbageCollectionAlgorithm
  twoSpaceCopying = GarbageCollectionAlgorithm
      "Two-space copying"
      garbageCollect
  
  newtype Gc a = Gc (ReaderT Heap (State GcState) a)
      deriving (Functor, Applicative, Monad)
  
  askHeap :: Gc Heap
  askHeap = Gc ask
  
  getGcState :: Gc GcState
  getGcState = Gc (lift get)
  
  putGcState :: GcState -> Gc ()
  putGcState s = Gc (lift (put s))
  
  execGc :: Gc a -> Heap -> GcState -> GcState
  execGc (Gc rsa) oldHeap gcState =
      let sa = runReaderT rsa oldHeap
          finalState = execState sa gcState
      in finalState
  
  data GcState = GcState
      { toHeap :: Heap
          -- ^ Heap of closures known to be alive.
  
      , forwards :: Map MemAddr MemAddr
          -- ^ Forward pointers to the new locations of already collected heap
          -- objects
  
      , toScavenge :: Seq MemAddr
          -- ^ Heap objects already evacuated, but not yet scavenged. Contains
          -- only objects that are also in the @toHeap@.
  
      , toEvacuate :: Seq MemAddr
          -- ^ Heap objects known to be alive, but not yet evacuated.
      } deriving (Eq, Ord, Show)
  
  garbageCollect :: StgState -> (Set MemAddr, Map MemAddr MemAddr, StgState)
  garbageCollect stgState@StgState
      { stgCode    = code
      , stgHeap    = heap
      , stgGlobals = globals
      , stgStack   = stack }
    = let rootAddrs = mconcat [addrs code, addrs stack, addrs globals]
          initialState = GcState
              { toHeap     = mempty
              , forwards   = mempty
              , toScavenge = mempty
              , toEvacuate = rootAddrs }
          finalState = execGc evacuateScavengeLoop heap initialState
      in case finalState of
          GcState {toHeap = heap', forwards = forwards'} ->
              let deadFormerAddrs
                    = let Heap old = heap
                      in M.keysSet old `S.difference` M.keysSet forwards'
  
                  forward addr = M.findWithDefault forwardErr addr forwards'
                  forwardErr = error "Invalid forward in GC; please report this as a bug"
  
                  removeIdentities = M.filterWithKey (/=)
  
                  stgState' = stgState
                      { stgCode    = updateAddrs forward code
                      , stgStack   = updateAddrs forward stack
                      , stgGlobals = updateAddrs forward globals
                      , stgHeap    = heap' }
              in (deadFormerAddrs, removeIdentities forwards', stgState')
  
  evacuateScavengeLoop :: Gc ()
  evacuateScavengeLoop = initialEvacuation >> scavengeLoop
  
  initialEvacuation :: Gc ()
  initialEvacuation = getAndClearToEvacuate >>= evacuateAll
    where
      getAndClearToEvacuate = do
          gcState <- getGcState
          putGcState (gcState{toEvacuate = mempty})
          pure (toEvacuate gcState)
      evacuateAll = traverse_ evacuate
  
  scavengeLoop :: Gc ()
  scavengeLoop = do
      scavengeNext <- getAndClearToScavenge
      if | Seq.null scavengeNext -> pure ()
         | otherwise -> do
             scavengeAddrs S.empty scavengeNext
             scavengeLoop
    where
      getAndClearToScavenge = do
          gcState <- getGcState
          putGcState (gcState{toScavenge = mempty})
          pure (toScavenge gcState)
      scavengeAddrs alreadyScavenged toAddrs = case Seq.viewl toAddrs of
          EmptyL -> pure ()
          addr :< rest
              | S.member addr alreadyScavenged -> scavengeAddrs alreadyScavenged rest
              | otherwise -> do
                  scavenge addr
                  scavengeAddrs (S.insert addr alreadyScavenged) rest
  
  data EvacuationStatus = NotEvacuatedYet | AlreadyEvacuatedTo MemAddr
  
  -- | Copy a closure from from-space to to-space, if it has not been evacuated
  -- previously. Copying to to-space involves a new allocation in to-space,
  -- registering the copied address to be scavenged, and creating a forwarding
  -- entry so that further evacuations can short-circuit.
  evacuate :: MemAddr -> Gc MemAddr
  evacuate = \fromAddr -> forwardingStatus fromAddr >>= \case
      AlreadyEvacuatedTo newAddr -> pure newAddr
      NotEvacuatedYet -> fmap (H.lookup fromAddr) askHeap >>= \case
          Nothing -> error "Tried collecting a non-existent memory address!\
                           \ Please report this as a bug."
          Just heapObject -> do
              newAddr <- copyIntoToSpace heapObject
              registerToBeScavenged newAddr
              createForward fromAddr newAddr
              pure newAddr
    where
      forwardingStatus :: MemAddr -> Gc EvacuationStatus
      forwardingStatus addr = do
          GcState { forwards = forw } <- getGcState
          pure (case M.lookup addr forw of
              Nothing -> NotEvacuatedYet
              Just newAddr -> AlreadyEvacuatedTo newAddr )
  
      copyIntoToSpace :: HeapObject -> Gc MemAddr
      copyIntoToSpace heapObject = do
          gcState <- getGcState
          let (addr', to') = H.alloc heapObject (toHeap gcState)
          putGcState gcState { toHeap = to' }
          pure addr'
  
      registerToBeScavenged :: MemAddr -> Gc ()
      registerToBeScavenged addr = do
          gcState@GcState { toScavenge = sc } <- getGcState
          putGcState gcState { toScavenge = sc |> addr }
  
      createForward :: MemAddr -> MemAddr -> Gc ()
      createForward from to = do
          gcState@GcState{forwards = forw} <- getGcState
          putGcState gcState { forwards = M.insert from to forw }
  
  -- | Find referenced addresses in a heap object, and overwrite them with their
  -- evacuated new addresses.
  scavenge :: MemAddr -> Gc ()
  scavenge = \scavengeAddr -> do
      scavengeHeapObject <- do
          GcState { toHeap = heap } <- getGcState
          pure (H.lookup scavengeAddr heap)
      case scavengeHeapObject of
          Nothing -> error "Scavenge error: address not found on to-heap\
                           \ Please report this as a bug."
          Just Blackhole{} -> pure mempty
          Just (HClosure (Closure lf frees)) -> do
              frees' <- evacuateContainedValues frees
              updateClosure scavengeAddr (Closure lf frees')
              registerForEvacuation [ addr | Addr addr <- frees' ]
    where
      evacuateContainedValues :: [Value] -> Gc [Value]
      evacuateContainedValues = traverse (\case
          Addr addr -> fmap (\x -> Addr x) (evacuate addr)
          i@PrimInt{} -> pure i )
  
      updateClosure :: MemAddr -> Closure -> Gc ()
      updateClosure addr closure = do
          gcState@GcState { toHeap = heap } <- getGcState
          let heap' = H.update (Mapping addr (HClosure closure)) heap
          putGcState gcState { toHeap = heap' }
  
      registerForEvacuation :: [MemAddr] -> Gc ()
      registerForEvacuation addresses = do
          gcState@GcState { toEvacuate = evac } <- getGcState
          putGcState gcState { toEvacuate = evac <> Seq.fromList addresses }
  