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.Monoid
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
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
, forwards :: Map MemAddr MemAddr
, toScavenge :: Seq MemAddr
, toEvacuate :: Seq MemAddr
} 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
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 }
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 }