content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  -- | Tri-state ("tri-colour") garbage collector.
  --
  -- * Not compacting: alive memory is not altered
  -- * Tracing
  module Stg.Machine.GarbageCollection.TriStateTracing (
      triStateTracing,
  ) where
  
  
  
  import           Data.Map      (Map)
  import qualified Data.Map      as M
  import           Data.Sequence (Seq)
  import           Data.Set      (Set)
  import qualified Data.Set      as S
  
  import Stg.Machine.GarbageCollection.Common
  import Stg.Machine.Types
  
  
  
  -- | Remove all unused addresses, without moving the others.
  triStateTracing :: GarbageCollectionAlgorithm
  triStateTracing = GarbageCollectionAlgorithm
      "Tri-state tracing"
      (insert2nd mempty . garbageCollect)
  
  insert2nd :: a -> (x, y) -> (x, a, y)
  insert2nd a (x,y) = (x,a,y)
  
  garbageCollect :: StgState -> (Set MemAddr, StgState)
  garbageCollect stgState@StgState
      { stgCode    = code
      , stgHeap    = heap
      , stgGlobals = globals
      , stgStack   = stack }
    = let GcState {aliveHeap = alive, oldHeap = Heap dead}
              = until everythingCollected gcStep start
          start = GcState
              { aliveHeap = mempty
              , oldHeap = heap
              , staged = (seqToSet . mconcat)
                  [addrs code, addrs globals, addrs stack] }
          stgState' = stgState { stgHeap = alive }
      in (M.keysSet dead, stgState')
  
  seqToSet :: Ord a => Seq a -> Set a
  seqToSet = foldMap S.singleton
  
  everythingCollected :: GcState -> Bool
  everythingCollected = noAlives
    where
      noAlives GcState {staged = alive} = S.null alive
  
  -- | Each closure is in one of three states: in the alive heap, staged for
  -- later rescue, or not even staged yet.
  data GcState = GcState
      { aliveHeap :: Heap
          -- ^ Heap of closures known to be alive.
          --   Has no overlap with the old heap.
  
      , staged :: Set MemAddr
          -- ^ Memory addresses known to be alive,
          --   but not yet rescued from the old heap.
  
      , oldHeap :: Heap
          -- ^ The old heap, containing both dead
          --   and not-yet-found alive closures.
      } deriving (Eq, Ord, Show)
  
  gcStep :: GcState -> GcState
  gcStep GcState
      { aliveHeap = oldAlive@(Heap alive)
      , staged    = stagedAddrs
      , oldHeap   = Heap oldRest }
    = GcState
      { aliveHeap = oldAlive <> Heap rescued
      , staged    = seqToSet (addrs rescued)
      , oldHeap   = Heap newRest }
    where
      rescued, newRest :: Map MemAddr HeapObject
      (rescued, newRest) = M.partitionWithKey isAlive oldRest
        where
          isAlive addr _closure = M.member addr alive
                               || S.member addr stagedAddrs
  