-----------------------------------------------------------------------------
-- |
-- Module      :  Windll
-- Copyright   :  (c) Tamar Christina 2009 - 2010
-- License     :  BSD3
-- 
-- Maintainer  :  tamar@zhox.com
-- Stability   :  experimental
-- Portability :  portable
--
-- A virtual heap manipulation function
--
-----------------------------------------------------------------------------

module WinDll.Debug.Heap where

import Numeric
import Data.Char

import WinDll.Debug.Records
import WinDll.Session.Debug
import WinDll.Utils.Feedback

import Data.Function
import Data.List

type Address = Int
type Size    = Int

-- | Virtual heap representation
data Heap = Blob {start :: Address 
                 ,stop  :: Address 
                 }
          | Heap {value :: Heap
                 ,next  :: Heap 
                 }
          | NoHeap
    deriving (Eq, Show)
    
asHex = \x -> showIntAtBase 16 intToDigit x ""

-- | Retreive the amount of items in the heap
sizeHeap :: Heap -> Int
sizeHeap NoHeap = 0
sizeHeap Blob{} = 1
sizeHeap heap   = sizeHeap (value heap) + sizeHeap (next heap)

-- | Get all the starting addresses
getStartingAddr :: Heap -> [Address]
getStartingAddr NoHeap     = []
getStartingAddr (Heap a b) = getStartingAddr a ++ getStartingAddr b
getStartingAddr x          = [start x]
    
-- | Add a new address to the heap    
addToHeap :: Heap -> Address -> Size -> Exec Heap
addToHeap heap st sz = do inform _detail $ "Mallocing " ++ show sz ++ " bytes starting at 0x" ++ asHex st ++ "."
                          return $ insertHeap heap (Blob st (st + sz))

-- | Grow the heap at the starting address by the amount given.
--   New size is oldsize + given size
expandHeap :: Heap -> Address -> Size -> Exec Heap
expandHeap heap start sz = do inform _detail $ "Expanding 0x" ++ asHex start ++ " with " ++ show sz ++ " bytes."
                              modifyHeap heap start mk
  where mk (Blob st os) = Blob st (os + sz)

-- | Shrink the heap at the starting address by the amount given.
--   New size is oldsize - given size
contractHeap :: Heap -> Address -> Size -> Exec Heap
contractHeap heap start sz = do inform _detail $ "Contracting 0x" ++ asHex start ++ " with " ++ show sz ++ " bytes."
                                modifyHeap heap start mk
  where mk (Blob st os) = Blob st (os - sz)

-- | Resize the address starting at the given point to the given size
resizeHeap :: Heap -> Address -> Size -> Exec Heap
resizeHeap heap start sz = do inform _detail $ "Resizing 0x" ++ asHex start ++ " to " ++ show sz ++ " bytes."
                              modifyHeap heap start mk
  where mk (Blob st os) = Blob st sz
  
-- | Merge two heaps together producing one larger heap
combineHeap :: Heap -> Heap -> Exec Heap
combineHeap hp1 hp2 = do inform _detail "Merging heaps..."
                         let heap = merge hp1 hp2
                         inform _detail "Merge done."
                         normalizeHeap heap
  where merge NoHeap       hp2      = hp2
        merge hp1          NoHeap   = hp1
        merge hp1          h@Blob{} = Heap h hp1
        merge h@Blob{}     hp2      = Heap h hp2        
        merge (Heap v1 n1) hp2      = merge n1 (Heap v1 hp2)
        
-- | Remove an entry from the heap
removeFromHeap :: Heap -> Address -> Exec Heap
removeFromHeap heap addr = do inform _detail $ "Removing pointer 0x" ++ asHex addr ++ " from heap."
                              heap <- modifyHeap heap addr (const NoHeap)
                              -- normalizeHeap heap
                              return heap

-- | Reduce the heap to the smallest possible heap
--   without merging bordering allocations. This might be added later
--   this function assumes the heap is sorted.
normalizeHeap :: Heap -> Exec Heap
normalizeHeap heap = let lst     = flatten heap 
                         entries = sortBy (compare `on` start) lst
                     in  return $ foldl' insertHeap NoHeap entries
  where flatten :: Heap -> [Heap]
        flatten b@Blob{}   = [b]
        flatten NoHeap     = []
        flatten (Heap a b) = flatten a ++ flatten b
        
-- | Insert an item into a sorted heap, keeping it sorted
insertHeap :: Heap -> Heap -> Heap
insertHeap NoHeap     heap       = heap
insertHeap heap       NoHeap     = heap
insertHeap hp1@Blob{} hp2@Blob{} = if start hp1 > stop hp2
                                  then Heap hp2 hp1
                                  else if start hp2 > stop hp1
                                          then Heap hp1 hp2
                                          else Blob (start hp1 `min` start hp2) (stop hp1 `max` stop hp2)
insertHeap heap1      heap2      = heap1 { next = insertHeap (next heap1) heap2 }

-- | Looks up the heap allocation starting at the given address.
--   If it's found, it's updated by applying the update function.
--   If it's not found the unmodified heap would be returned.
modifyHeap :: Heap -> Address -> (Heap -> Heap) -> Exec Heap
modifyHeap heap start f = return $ modify heap
  where modify h@(Blob st _) | st == start = f h
        modify (Heap val nxt)              = let val' = modify val
                                                 nxt' = modify nxt
                                             in if val' == val
                                                   then Heap val  nxt'
                                                   else Heap val' nxt
        modify x                           = x