Skip to content

Instantly share code, notes, and snippets.

@glguy
Last active August 29, 2015 14:18
Show Gist options
  • Save glguy/2a613e0089342995be6b to your computer and use it in GitHub Desktop.
Save glguy/2a613e0089342995be6b to your computer and use it in GitHub Desktop.
mark and sweep
module MarkSweep where
import Data.Array.IO
import Data.Word
import Data.Bits
import Data.Bits.Lens
import Data.Foldable (traverse_)
import Control.Lens
import Control.Monad
------------------------------------------------------------------------
-- Heaps
------------------------------------------------------------------------
{- Heap layout: contiguous sequence of objects
- Object layout:
- bit 0: mark bit
- bit 1: allocated bit
- bit 31-2: size (maximum allocation is therefore 2^30-1
- -}
newtype Heap = Heap (IOUArray Int Word32)
readHeap ::
Heap {- ^ heap -} ->
Int {- ^ read address -} ->
IO HeapElem
readHeap (Heap a) i = fmap HeapElem (readArray a i)
writeHeap ::
Heap {- ^ heap -} ->
Int {- ^ write address -} ->
HeapElem {- ^ new value -} ->
IO ()
writeHeap (Heap a) i (HeapElem e) = writeArray a i e
-- | Determine range of valid indexes in a heap
heapBounds :: Heap -> IO (Int,Int)
heapBounds (Heap a) = getBounds a
-- | Construct an empty heap with a single free block
-- of the given size
initialHeap :: Int -> IO Heap
initialHeap sz =
do heap <- fmap Heap (newArray (0,sz-1) 0)
writeHeap heap 0 (mkFreeBlock sz)
return heap
-- | Attempt to allocate a new block in the heap of the given size.
allocate ::
Heap {- ^ heap -} ->
Int {- ^ allocation size -} ->
IO Int {- ^ index of allocation -}
allocate heap sz =
do b <- heapBounds heap
next b 0
where
next :: (Int,Int) -> Int -> IO Int
next b i
| inRange b i = attemptAllocation b i =<< readHeap heap i
| otherwise = fail "virtual heap exhausted"
attemptAllocation b i e
| e^.allocated || esz < sz = next b (i + e^.elemSize)
| otherwise =
do -- mark unused portion of this block as free
when (sz < esz)
(writeHeap heap (i+sz) (mkFreeBlock (esz-sz)))
let e' = set allocated True
$ mkFreeBlock sz
writeHeap heap i e'
return i
where
esz = e^.elemSize
-- | Walk through heap from beginning to end
-- combining subsequent free blocks.
coalesceHeap :: Heap -> IO ()
coalesceHeap h =
do (lo,hi) <- heapBounds h
let next i =
when (i <= hi) $
do e <- readHeap h i
let i' = i + view elemSize e
if view allocated e || i' >= hi then next i'
else do e' <- readHeap h i'
if view allocated e'
then next (i' + view elemSize e')
else do writeHeap h i (mkFreeBlock (view elemSize e + view elemSize e'))
next i -- retry
next lo
describeHeap :: Heap -> IO ()
describeHeap h =
do (lo,hi) <- heapBounds h
let next i =
when (i <= hi) $
do e <- readHeap h i
print (i,view allocated e, view marked e, view elemSize e)
next (i + view elemSize e)
next lo
------------------------------------------------------------------------
-- Heap elements
------------------------------------------------------------------------
newtype HeapElem = HeapElem { heapElemRep :: Word32 }
_HeapElem :: Iso' HeapElem Word32
_HeapElem = iso heapElemRep HeapElem
mkFreeBlock :: Int -> HeapElem
mkFreeBlock sz
| 0 <= sz && sz < 2^30-1 = HeapElem (fromIntegral sz `shiftL` 2)
| otherwise = error "mkFreeBlock: size out of range"
marked :: Lens' HeapElem Bool
marked = _HeapElem . bitAt 0
allocated :: Lens' HeapElem Bool
allocated = _HeapElem . bitAt 1
elemSize :: Lens' HeapElem Int
elemSize = _HeapElem
. lens (\s -> fromIntegral (shiftR s 2))
(\s b -> shiftL (fromIntegral b) 2 .|. (0x3 .&. s))
------------------------------------------------------------------------
-- Object layouts
------------------------------------------------------------------------
data ObjectType = SumType | ProductType
data ObjectDescription = ObjectDescription
{ objectType :: ObjectType
, objectFields :: [FieldType]
}
data FieldType
= IntField
| ObjectField ObjectDescription
allocateSum ::
Heap ->
Int {- ^ alternative tag -} ->
Word32 {- ^ value of alternative -} ->
IO Int {- ^ pointer to allocated and initialized block -}
allocateSum h alt v =
do p <- allocate h 3
writeHeap h (p+1) (HeapElem (fromIntegral alt))
writeHeap h (p+2) (HeapElem v)
return p
allocateProduct ::
Heap ->
[Word32] {- ^ list of fields in product -} ->
IO Int {- ^ pointer to allocated and initialized block -}
allocateProduct h vs =
do p <- allocate h (1+length vs)
zipWithM_ (\i e -> writeHeap h i (HeapElem e))
[p+1, p+2..]
vs
return p
------------------------------------------------------------------------
-- Sample object descriptions
------------------------------------------------------------------------
intObject :: ObjectDescription
intObject = ObjectDescription
{ objectType = ProductType
, objectFields = [IntField]
}
mkInt ::
Heap {- ^ allocation heap -} ->
Int {- ^ int value -} ->
IO Int {- ^ returns pointer to boxed int value -}
mkInt h v = allocateProduct h [fromIntegral v]
pairObject :: ObjectDescription -> ObjectDescription -> ObjectDescription
pairObject a b = ObjectDescription
{ objectType = ProductType
, objectFields = [ObjectField a, ObjectField b]
}
mkPair ::
Heap ->
Int {- ^ fst pointer -} ->
Int {- ^ snd pointer -} ->
IO Int
mkPair h x1 x2 = allocateProduct h [fromIntegral x1, fromIntegral x2]
unitObject :: ObjectDescription
unitObject = ObjectDescription
{ objectType = ProductType
, objectFields = []
}
mkUnit :: Heap -> IO Int
mkUnit h = allocateProduct h []
maybeObject :: ObjectDescription -> ObjectDescription
maybeObject a = ObjectDescription
{ objectType = SumType
, objectFields = [ObjectField unitObject, ObjectField a]
}
mkNothing :: Heap -> IO Int
mkNothing h =
do p <- mkUnit h
allocateSum h 0 (fromIntegral p)
mkJust :: Heap -> Int -> IO Int
mkJust h v = allocateSum h 1 (fromIntegral v)
listObject :: ObjectDescription -> ObjectDescription
listObject a = maybeObject (pairObject a (listObject a))
mkNil :: Heap -> IO Int
mkNil = mkNothing
mkCons ::
Heap ->
Int {- ^ pointer to head of list -} ->
Int {- ^ pointer to tail of list -} ->
IO Int {- ^ pointer to list -}
mkCons h x xs =
do p <- mkPair h x xs
allocateSum h 1 (fromIntegral p)
------------------------------------------------------------------------
-- Mark and sweep GC
------------------------------------------------------------------------
mark ::
Heap {- ^ heap to mark -} ->
ObjectDescription {- ^ description of the object at the address -} ->
Int {- ^ address to start marking -} ->
IO ()
mark h obj i =
do e <- readHeap h i
writeHeap h i (set marked True e)
case objectType obj of
SumType -> markSum h (objectFields obj) (i+1)
ProductType -> markProduct h (objectFields obj) (i+1)
markSum ::
Heap {- ^ heap to mark -} ->
[FieldType] {- ^ description of the possible field types -} ->
Int {- ^ address of sum type index -} ->
IO ()
markSum h alts i =
do altNum <- fmap (views _HeapElem fromIntegral) (readHeap h i)
case preview (ix altNum) alts of
Nothing -> fail ("Invalid sum type at " ++ show i)
Just alt -> markField h alt (i+1)
markProduct ::
Heap {- ^ heap to mark -} ->
[FieldType] {- ^ description of the sequentially stored fields -} ->
Int {- ^ address of the first field -} ->
IO ()
markProduct h fields i = zipWithM_ (markField h) fields [i,i+1..]
markField ::
Heap {- ^ heap to mark -} ->
FieldType {- ^ description of this field -} ->
Int {- ^ address of this field -} ->
IO ()
markField _ IntField _ = return ()
markField h (ObjectField obj) i =
do e <- readHeap h i
let i' = views _HeapElem fromIntegral e
mark h obj i'
-- | Walk through heap from beginning to end deallocating any
-- unmarked but allocated region.
sweep :: Heap -> IO ()
sweep h =
do (lo,hi) <- heapBounds h
let next i = when (i <= hi) $
do e <- readHeap h i
if view allocated e
then if view marked e
then do writeHeap h i (set marked False e)
next (i + view elemSize e)
else do writeHeap h i (set allocated False e)
next i
else next (i + view elemSize e)
next lo
collectGarbage ::
Heap {- ^ heap to gc -} ->
[(ObjectDescription,Int)] {- ^ live root types and addresses -} ->
IO ()
collectGarbage h roots =
do traverse_ (\(obj,i) -> mark h obj i) roots
sweep h
coalesceHeap h
------------------------------------------------------------------------
-- Test case
------------------------------------------------------------------------
demo :: IO ()
demo = do
h <- initialHeap 100
one <- mkInt h 1
two <- mkInt h 2
_three <- mkInt h 3
nil <- mkNil h
x3 <- mkCons h one nil
x2 <- mkCons h two x3
x1 <- mkCons h one x2
collectGarbage h [(listObject intObject, x1)]
putStrLn "After GC"
describeHeap h
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment