-- GRIN-like backend for Yhc Core.
-- Dead code elimination.

module Yhc.Core.GRIN.ElimDead where

import Yhc.Core.GRIN.Type
import Yhc.Core.GRIN.FindType
import Yhc.Core.GRIN.HeapPointsTo
import Yhc.Core.GRIN.SubstVars

import Data.List
import qualified Data.Set as S
import qualified Data.Map as M

-- Eliminate duplicate bindings. Scan each block from the beginning
-- accumuilating a set of used variables. If the current binding
-- binds to a variable already in the set, remove the binding.

elimDupBindings :: GFunc -> GFunc

elimDupBindings gf@(GFunc {gFuncBody = gb}) = gf {gFuncBody = f gb} where
  seen = S.fromList (gFuncArgs gf)
  f = filter (/= GNop) . snd . elimDB where
    elimDB = elimWith seen
    elimWith s = mapAccumL edbs s
    edbs us   (GSimple (GInline ib)) =
      let el = elimWith us ib
          newblk = filter (/= GNop) $ snd el
      in  (us, GSimple (GInline newblk))
    edbs us   (GCase cv blks) =
      let mblks = map (elimWith us . snd) blks
          pats = map fst blks
          newblks = map (filter (/= GNop) . snd) mblks
          newcase = GCase cv (zip pats newblks)
      in  (us, newcase)
    edbs us   (GBind (GInline ib) bv) =
      let el = elimWith us ib
          newblk = filter (/= GNop) $ snd el
          newset = S.insert (unGVar bv) us
      in  case S.member (unGVar bv) us of
            False -> (newset, GBind (GInline newblk) bv)
            True  -> (us, GNop)
    edbs us x@(GBind sx bv) =
      let newset = S.insert (unGVar bv) us
      in  case S.member (unGVar bv) us of
            False -> (newset, x)
            True  -> (us, GNop)
    edbs us z = (us, z)

-- Eliminate dead bindings. Scan each block from bottom-up (from the end)
-- accumulating a set of used variables (that is, appearing within
-- a GSexpr inside GBind). If the current binding binds to a variable
-- that does not appear on this set, remove the binding as it will
-- not be used further down the code. This is a function-level optimization.
-- If a CASE is the last element of a block of bindings, the function collects
-- used variables from all its branches.


elimDeadBindings :: GFunc -> GFunc

elimDeadBindings gf@(GFunc {gFuncBody = gb}) = gf {gFuncBody = f gb} where
  f = filter (/= GNop) . snd . elimDB where
    elimDB = mapAccumR edbs S.empty
    edbs us   (GSimple (GInline ib)) =
      let el = elimDB ib
          newblk = filter (/= GNop) $ snd el
          newset = S.union us (fst el)
      in  (newset, GSimple (GInline newblk))
    edbs us x@(GSimple sx) = (S.union us (useVars sx), x)
    edbs us   (GCase cv blks) = 
      let mblks = map (elimDB . snd) blks
          pats = map fst blks
          newblks = map (filter (/= GNop) . snd) mblks
          newset = S.unions (S.singleton (unGVar cv) : us : map fst mblks)
          newcase = GCase cv (zip pats newblks)
      in  (newset, newcase)
    edbs us   (GBind (GInline ib) bv) = 
      let el = elimDB ib
          newblk = filter (/= GNop) $ snd el
          newset = S.union us (fst el)
      in  case S.member (unGVar bv) us of
            True -> (newset, GBind (GInline newblk) bv)
            False -> (us, GNop)
    edbs us x@(GBind sx bv) = 
      let newset = S.union us (useVars sx)
      in  case S.member (unGVar bv) us of
            True -> (newset, x)
            False -> (us, GNop)
    edbs us z = (us, z)

-- Apply elimDeadBindings to each finction in GRIN

elimAllDeadBindings :: GRIN -> GRIN

elimAllDeadBindings g =
  g {gFuncs = map (elimDeadBindings . elimDupBindings) $ gFuncs (elimEquivVars g)}

-- Eliminate variables equivalent to others just by looking at the heap map.

elimEquivVars :: GRIN -> GRIN

elimEquivVars g =
  let hm = heapPointsTo g
      eqvar ((HVar v1), [EquivTo v2]) = [(v1, v2)]
      eqvar (_, _) = []
      eqm = M.fromList $ concat $ map eqvar $ M.toList hm
  in  substGRIN eqm g

-- Eliminate dead functions by analysing the heap map for calls
-- to functions vs. declared functions.

elimDeadFuncs :: [GName] -> HeapMap -> GRIN -> GRIN

elimDeadFuncs roots seed g =
  let hm = heapPointsTo g `M.union` seed
      ccall [CallTo f] = [f]
      ccall _ = []
      called = S.fromList $ roots ++ (concat $ map ccall $ M.elems hm)
  in  g {gFuncs = filter (flip S.member called . gFuncName) (gFuncs g)}

-- Eliminate dictionary fields that are not loaded.

elimDictFields :: HeapMap -> GRIN -> GRIN

elimDictFields seed g =
  let hm = heapPointsTo g `M.union` seed
      elts = S.fromList $ concat $ map elt $ M.elems $ M.filter eltof hm
      eltof [ElemOfT _ _ _] = True
      eltof _ = False
      elt [ElemOfT v t n] | not (null t) && n /= 0 = [(t, n)]
      elt _ = []
      mfun (GBind (GStore t@(GTagged _ _)) bv) = GBind (GStore (efld t)) bv
      mfun (GSimple (GUnit t@(GTagged _ _))) = GSimple (GUnit (efld t))
      mfun ex = ex
      efld (GTagged t vs) = GTagged t (zipWith (ff t) vs [1 .. ])
      ff t v n | S.member (t, n) elts = v
      ff _ v _ = case M.lookup (HVar (unVar v)) hm of
        Just [HasValue (GTag ('K':'@':_))] -> GVar "n0"
        Just [HasValue (GTagged ('F':'@':_) _)] -> GVar "n0"
        Just [HasValue (GTagged ('P':'@':_) _)] -> GVar "n0"
        Just _ -> v
        Nothing -> GVar "n0"
      mg f | gFuncName f `elem` ["GRIN;eval", "GRIN;apply"] = f
      mg f = f {gFuncBody = mapOverBlock mfun (gFuncBody f)}
  in  elimAllDeadBindings $ g {gFuncs = map mg (gFuncs g)}