--
-- TreeToGraph.hs
--
-- Copyright (c) 2007, 2008 Antiope Associates LLC, all rights reserved.
--

module TreeToGraph (
        nonzeroStacks,
        treeToGraph
) where

import Data.IntMap as IM
import Data.List as L

import ParseProfile


-- | nonzeroStacks is a utility function that take the IntMap of all stacks
--   and returns an IntMap containing only those stacks that have some
--   cost associated with them (i.e., are called at least once or allocate
--   some memory).
--
--   There is one corner case: the root stack should always be "nonzero" even
--   if no costs are attibuted to it.  Deleting it would break the tree to graph
--   conversion.
--
nonzeroStacks :: Profile -> IntMap CostCenterStack
nonzeroStacks p = let
        ss  = stacks p
        nzs = intersection ss (IM.filter (\s -> (reportCount s) /= 0 ||
                                                (reportTicks s) /= 0 ||
                                                (reportAlloc s) /= 0) (reports p))
        in
          IM.insert 1 (ss IM.! 1) nzs


sortCCID :: Profile -> (Int, CostCenterStack) -> (Int, CostCenterStack) -> Ordering
sortCCID _ (_, cc1) (_, cc2) = let
        id1 = stackCCID cc1
        id2 = stackCCID cc2
        in
          if id1 == id2
              then EQ else if id1 > id2
                               then GT
                               else LT


groupCCID :: Profile -> (Int, CostCenterStack) -> (Int, CostCenterStack) -> Bool
groupCCID _ (_, cc1) (_, cc2) = stackCCID cc1 == stackCCID cc2


collapseCosts :: Profile -> [ Int ] -> CostCenterReport
collapseCosts p l = foldl' addCosts CostCenterReport {reportCount = 0, reportTicks = 0, reportAlloc = 0} l
        where
          addCosts :: CostCenterReport -> Int -> CostCenterReport
          addCosts c n = CostCenterReport { reportCount = reportCount c + reportCount ((reports p) IM.! n),
                                            reportTicks = reportTicks c + reportTicks ((reports p) IM.! n),
                                            reportAlloc = reportAlloc c + reportAlloc ((reports p) IM.! n) }


collapseCosts' :: Profile -> [ (Int, CostCenterStack) ] -> (Int, CostCenterReport)
collapseCosts' p l = let
        ns = fst (unzip l)
        cs = collapseCosts p ns
        in
          (head ns, cs)


collapseStacks :: [ (Int, CostCenterStack) ] -> (Int, CostCenterStack)
collapseStacks l = let
        (ns, cs) = unzip l
        ps = concatMap parentStack cs
        in
          (head ns, CostCenterStack { stackCCID   = stackCCID (head cs),
                                      parentStack = ps })


newParentMap :: [[ (Int, CostCenterStack) ]] -> IntMap Int
newParentMap sl = let
        equivClasses :: [[ (Int, CostCenterStack) ]] -> [[ Int ]]
        equivClasses ls = L.map (L.map fst) ls

        mkMap :: [[ Int ]] -> [ (Int, Int) ]
        mkMap ls = concatMap oneMap ls 

        oneMap :: [ Int ] -> [ (Int, Int) ]
        oneMap l = let
                       val = head l
                   in
                     L.map (\e -> (e, val)) l

        in
          fromList (mkMap (equivClasses sl))


reparent :: IntMap Int -> (Int, CostCenterStack) -> (Int, CostCenterStack)
reparent m (n, cs) = let
        np = nub $ L.map (m IM.!) (parentStack cs)
        in
          (n, CostCenterStack { stackCCID   = stackCCID cs,
                                parentStack = np })
        

-- | treeToGraph converts the profile representing the call tree
--   of the program into a call graph. In the call tree, a function
--   appears more than once if it is called by different parents.
--   This has the advantage of making it clear to which parent 
--   inherits a particular center's costs.  It has the
--   drawback of obscuring the structure of the program.
--
--   The treeToGraph function collapses the call tree to a call graph
--   adding up all of a center's cost.  The new Profile can be converted
--   to dot format just as the original call tree Profile was.
--
treeToGraph :: Profile -> Profile
treeToGraph p = let
        st  = nonzeroStacks p
        st' = groupBy (groupCCID p) (sortBy (sortCCID p) (toList st))

        -- The list st' now has the stacks grouped by identity.
        -- The next thing to do is to collapse the corresponding costs
        -- and build a new IntMap of the costs:

        reports' = fromList (L.map (collapseCosts' p) st')

        -- Before building the new IntMap of the stacks, we need
        -- to change the parent stack lists to repoint any stacks
        -- that will be deleted to the collapsed stack.

        newParents = reparent (newParentMap st')
        
        -- build a new IntMap of the stacks:

        stacks'  = fromList (L.map (newParents . collapseStacks) st')

        in
          Profile { timestamp    = timestamp p,
                    tickInterval = tickInterval p,
                    profileTicks = profileTicks p,
                    centers      = centers p,
                    stacks       = stacks',
                    reports      = reports' }