parse.hs

{-# LANGUAGE BangPatterns #-}
module Main where

import qualified Data.ByteString.Lazy.Char8 as BS
import qualified Data.IntMap as M
import Data.Maybe
import Data.Bits
import Data.Char
import System.Environment
import Text.Printf
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.STM.TBMChan
import Control.Monad

type AllocMap = M.IntMap (Int, [Int])

readHex :: BS.ByteString -> Int
readHex = BS.foldl' digit 0
    where
    digit i c = shiftL i 4 + toVal c
    toVal v | isDigit v = ord v - ord '0'
    toVal v = ord v - ord 'a' + 10

printStack :: (Int, [Int]) -> String
printStack (p, stack) = printf "%d:" p ++ concatMap (printf " %x") stack

counter :: TBMChan (Char, Int, [Int]) -> TMVar (Int, Int, AllocMap) -> Int -> Int -> AllocMap -> Int -> Int -> AllocMap -> IO ()
counter chunk result !maxb !maxc maxallocs !b !c !allocs = do
    val <- atomically $ readTBMChan chunk
    case val of
        Nothing -> atomically $ putTMVar result (maxb, maxc, maxallocs)
        Just ('f', addr, []) -> if addr == 0 then counter chunk result maxb maxc maxallocs b c allocs
                                                 else let size = fst $ fromJust $ M.lookup addr allocs
                                                          in counter chunk result maxb maxc maxallocs (b - size) (c - 1) (M.delete addr allocs)
        Just ('a', addr, size:stack) -> let newallocs = M.insert addr (size, stack) allocs
                                            newmaxc = max maxc (c + 1)
                                            in if maxb < b + size then counter chunk result (b + size) newmaxc newallocs (b + size) (c + 1) newallocs
                                                                  else counter chunk result maxb newmaxc maxallocs (b + size) (c + 1) newallocs

main :: IO()
main = do
    filename <- head `fmap` getArgs
    src <- BS.readFile filename
    chunk <- atomically $ newTBMChan 10
    result <- atomically $ newEmptyTMVar
    _ <- forkIO $ counter chunk result 0 0 M.empty 0 0 M.empty

    forM_ (BS.lines src) $ \l -> do
            let (op:addr:params) = BS.words l
            let !opch = BS.head op
            let !addrval = readHex addr
            let !paramval = map (id $! readHex) params
            atomically $ writeTBMChan chunk (opch, addrval, paramval)
    atomically $ closeTBMChan chunk
    (maxbytes, maxchunks, allocs) <- atomically $ takeTMVar result
    print maxbytes
    print maxchunks
    let total = sum $ map fst (M.elems allocs)
    print total
--    mapM_ (printf "%s\n" . printStack) (M.elems allocs)