Add mark sweep in hs

diff --git a/mark-sweep/MarkSweep.hs b/mark-sweep/MarkSweep.hs
new file mode 100644 (file)
index 0000000..0bb07f0
--- /dev/null
+++ b/mark-sweep/MarkSweep.hs
@@ -0,0 +1,116 @@
+module MarkSweep where
+
+import Debug.Trace
+
+-- run as ghci MarkSweep.hs -package unordered-containers
+
+data Value = IntVal Int
+  deriving (Eq, Show)
+
+data Object = Object {
+  val :: Value,
+  fields :: [ObjectPtr],
+  marked :: Bool
+} deriving (Eq, Show)
+
+data ObjectPtr = ObjectPtr Int | NullPtr
+  deriving (Eq, Show)
+
+increment :: ObjectPtr -> ObjectPtr
+increment (ObjectPtr p) = ObjectPtr $ p + 1
+increment NullPtr = NullPtr
+
+newtype Heap = Heap { heapObjects :: [Maybe Object] }
+  deriving (Eq, Show)
+
+newHeap :: Heap
+newHeap = Heap [Nothing | x <- [1..8]]
+
+isMarked :: ObjectPtr -> Heap -> Bool
+isMarked NullPtr _ = False
+isMarked (ObjectPtr p) h = Just True == (marked <$> (heapObjects h !! p))
+
+setMarked :: ObjectPtr -> Heap -> Heap
+setMarked NullPtr h = h
+setMarked (ObjectPtr p) (Heap hs) = Heap $ reverse $ foldl step [] $ zip [0..] hs
+  where
+    step res (idx, Nothing) = Nothing : res
+    step res (idx, Just obj)
+      | p == idx = (Just $ obj { marked = True }) : res
+      | otherwise = (Just obj) : res
+
+newtype Roots = Roots { rootsObjects :: [ObjectPtr]}
+  deriving (Eq, Show)
+
+data MarkSweepGC = MarkSweepGC {
+  roots :: Roots,
+  heap :: Heap
+} deriving (Eq, Show)
+
+allocate :: Value -> Heap -> Either String (ObjectPtr, Heap) 
+allocate v (Heap objs) = case foldl step (False, -1, []) objs of
+  (True, idx, res) -> Right $ (ObjectPtr idx, Heap $ reverse res)
+  _ -> Left "Out of memory"
+  where
+    step (True, idx, res) obj = (True, idx, obj : res)
+    step (_, idx, res) Nothing = (True, idx + 1, (Just $ Object v [] False) : res) 
+    step (ok, idx, res) obj = (ok, idx + 1, obj : res)
+
+mark :: [ObjectPtr] -> Heap -> Heap
+mark [] h = h
+mark ((ObjectPtr p):rest) h = case heapObjects h !! p of
+  Just obj ->
+    let
+      step acc NullPtr = acc
+      step (wl, heap) child
+        | not (isMarked child h) = (child : wl, setMarked child heap)
+      step acc _ = acc
+      
+      (workList, newH) = foldl step ([], h) $ fields obj
+    in mark (workList ++ rest) newH
+  Nothing -> mark rest h
+
+markFromRoots :: Heap -> Roots -> (Heap, Roots)
+markFromRoots h rs =
+  let newHeap = foldl step h $ rootsObjects rs
+  in (newHeap, rs)
+  where
+    step heap ptr
+      | not (isMarked ptr heap) =
+        let newH = setMarked ptr heap
+        in mark [ptr] newH
+    step heap ptr = heap
+
+collect :: Heap -> Roots -> (Heap, Roots)
+collect h r =
+  let (h', r') = markFromRoots h r
+  in (sweep h', r')
+
+sweep :: Heap -> Heap
+sweep (Heap hs) = Heap $ map step hs
+  where
+    step Nothing = Nothing
+    step (Just o) = if marked o then Just (o { marked = False }) else Nothing
+
+newObject :: Value -> MarkSweepGC -> Either String (ObjectPtr, MarkSweepGC)
+newObject v (MarkSweepGC roots heap) =
+  case allocate v heap of 
+    Right (ptr, heap') -> Right (ptr, MarkSweepGC roots heap') 
+    Left s ->
+      let (heap', roots') = collect heap roots
+      in case allocate v heap' of
+        Right (ptr, heap'') -> Right (ptr, MarkSweepGC roots' heap'')
+        Left s -> Left s
+
+main :: IO ()
+main = do
+  let gc = MarkSweepGC (Roots []) newHeap 
+  let
+    res = do
+        (ptr, gc') <- newObject (IntVal 3) gc
+        (ptr2, MarkSweepGC (Roots roots) h) <- newObject (IntVal 5) gc'
+
+        let (r'', h'') = collect h (Roots $ ptr2 : roots)
+        let newGc = MarkSweepGC h'' r''
+        return ([ptr, ptr2], newGc)
+  print res