content: -- | The STG heap maps memory addresses to closures.
  module Stg.Machine.Heap (
      -- * Info
      size,
  
      -- * Management
      lookup,
      update,
      updateMany,
      alloc,
      allocMany,
  ) where
  
  
  
  import qualified Data.List   as L
  import qualified Data.Map    as M
  import           Data.Monoid
  import           Prelude     hiding (lookup)
  
  import Stg.Machine.Types
  
  
  
  -- | Current number of elements in a heap.
  size :: Heap -> Int
  size (Heap heap) = M.size heap
  
  -- | Look up a value on the heap.
  lookup :: MemAddr -> Heap -> Maybe HeapObject
  lookup addr (Heap heap) = M.lookup addr heap
  
  -- | Update a value on the heap.
  update :: Mapping MemAddr HeapObject -> Heap -> Heap
  update (Mapping addr obj) (Heap h) = Heap (M.adjust (const obj) addr h)
  
  -- | Update many values on the heap.
  updateMany :: [Mapping MemAddr HeapObject] -> Heap -> Heap
  updateMany mappings heap =
      L.foldl' (\h mapping -> update mapping h) heap mappings
  
  -- | Store a value in the heap at an unused address.
  alloc :: HeapObject -> Heap -> (MemAddr, Heap)
  alloc lambdaForm heap = (addr, heap')
    where
      ([addr], heap') = allocMany [lambdaForm] heap
  
  -- | Store many values in the heap at unused addresses, and return them
  -- in input order.
  allocMany :: [HeapObject] -> Heap -> ([MemAddr], Heap)
  allocMany heapObjects (Heap heap) = (addrs, heap')
    where
      smallestUnusedAddr = case M.maxViewWithKey heap of
          Nothing -> 0
          Just ((MemAddr maxAddr, _), _) -> maxAddr+1
      addrs = zipWith (\addr _obj -> MemAddr addr)
                      [smallestUnusedAddr ..]
                      heapObjects
      heap' = Heap (heap <> M.fromList (zip addrs heapObjects))
  