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