content: {-# LANGUAGE NoImplicitPrelude #-}
  
  -- |
  -- Dead code elimination for `CoreFn`.
  module Language.PureScript.DCE.CoreFn
    ( runDeadCodeElimination
    , runBindDeadCodeElimination
    ) where
  
  import           Prelude hiding (mod)
  import           Control.Arrow ((***))
  import Control.Monad ( guard )
  import Data.Graph ( graphFromEdges, reachable, Vertex )
  import           Data.Foldable (foldl', foldr')
  import           Data.List (groupBy, sortBy)
  import qualified Data.Map.Strict as M
  import           Data.Maybe (catMaybes, mapMaybe)
  import qualified Data.Set as S
  import Language.PureScript.CoreFn
      ( Literal(ObjectLiteral, ArrayLiteral),
        Ann,
        everywhereOnValues,
        Binder(VarBinder, LiteralBinder, ConstructorBinder, NamedBinder),
        Bind(..),
        CaseAlternative(CaseAlternative),
        Expr(..),
        Module(Module, moduleImports, moduleExports, moduleForeign,
               moduleDecls, moduleName, moduleReExports) )
  import           Language.PureScript.DCE.Utils (bindIdents, unBind)
  import Language.PureScript.Names
      ( getQual,
        isQualified,
        mkQualified,
        Ident(Ident),
        ModuleName,
        ProperName(runProperName),
        Qualified(..),
        QualifiedBy(..) )
  
  type Key = Qualified Ident
  
  
  data DCEVertex
    = BindVertex (Bind Ann)
    | ForeignVertex (Qualified Ident)
    | ReExportedVertex (Qualified Ident)
  
  
  -- | Dead code elimination of a list of modules module
  --
  runDeadCodeElimination
    :: [Qualified Ident]
    -- ^ entry points used to build the graph of
    -- dependencies across module boundaries
    -> [Module Ann]
    -- ^ modules to dce
    -> [Module Ann]
    -- ^ dead code eliminated modules
  runDeadCodeElimination entryPoints modules = uncurry runModuleDeadCodeElimination `map` reachableInModule
    where
      -- DCE of a single module.
      runModuleDeadCodeElimination
        :: [(DCEVertex, Key, [Key])]
        -- list of qualified names that need to be preserved
        -> Module Ann
        -> Module Ann
      runModuleDeadCodeElimination vs mod@Module{ moduleDecls
                          , moduleExports
                          , moduleReExports
                          , moduleImports
                          , moduleName
                          , moduleForeign
                          } =
        let
            -- | filter declarations preserving the order
            moduleDecls' :: [Bind Ann]
            moduleDecls' = runBindDeadCodeElimination `map` filter filterByIdents moduleDecls
              where
              declIdents :: [Ident]
              declIdents = concatMap toIdents vs
  
              toIdents :: (DCEVertex, Key, [Key]) -> [Ident]
              toIdents (BindVertex b, _, _) = bindIdents b
              toIdents _                    = []
  
              filterByIdents :: Bind Ann -> Bool
              filterByIdents = any (`elem` declIdents) . bindIdents
  
            idents :: [Ident]
            idents = concatMap bindIdents moduleDecls'
  
            moduleExports' :: [Ident]
            moduleExports' =
              filter (`elem` (idents ++ moduleForeign')) moduleExports
  
            moduleReExports' :: M.Map ModuleName [Ident]
            moduleReExports' =
              filter (`elem` rexpIdents) <$> moduleReExports
              where
                toRexpIdents :: (DCEVertex, Key, [Key]) -> [Ident]
                toRexpIdents (ReExportedVertex (Qualified _ i), _, _) = [i]
                toRexpIdents _ = []
  
                rexpIdents :: [Ident]
                rexpIdents = concatMap toRexpIdents vs
  
            mods :: [ModuleName]
            mods = mapMaybe getQual (concatMap (\(_, _, ks) -> ks) vs)
  
            moduleImports' :: [(Ann, ModuleName)]
            moduleImports' = filter ((`elem` mods) . snd) moduleImports
  
            moduleForeign' :: [Ident]
            moduleForeign' = filter
                ((`S.member` reachableSet) . Qualified (ByModuleName moduleName))
                moduleForeign
              where
                reachableSet = foldr'
                  (\(_, k, ks) s -> S.insert k s `S.union` S.fromList ks)
                  S.empty vs
  
        in mod { moduleImports = moduleImports'
               , moduleExports = moduleExports'
               , moduleReExports = moduleReExports'
               , moduleForeign = moduleForeign'
               , moduleDecls   = moduleDecls'
               }
  
      (graph, keyForVertex, vertexForKey) = graphFromEdges verts
  
      -- | The Vertex set.
      verts :: [(DCEVertex, Key, [Key])]
      verts = do
          Module _ _ mn _ _ _ rexp mf ds <- modules
          concatMap (toVertices mn) ds -- Module local bindings
            ++ ((\q -> (ForeignVertex q, q, [])) . flip mkQualified mn) `map` mf -- Foreign bindings
            ++ reExportedVertices mn rexp -- Re-exported bindings
        where
        toVertices :: ModuleName -> Bind Ann -> [(DCEVertex, Key, [Key])]
        toVertices mn b@(NonRec _ i e) =
          [(BindVertex b, mkQualified i mn, deps e)]
        toVertices mn b@(Rec bs) =
          let ks :: [(Key, [Key])]
              ks = map (\((_, i), e) -> (mkQualified i mn, deps e)) bs
          in map (\(k, ks') -> (BindVertex b, k, map fst ks ++ ks')) ks
  
        reExportedVertices :: ModuleName -> M.Map ModuleName [Ident] -> [(DCEVertex, Key, [Key])]
        reExportedVertices parent rexp =
          filter (isReExported rexp) modules >>= rexpsFor
          where
            rexpsFor :: Module Ann -> [(DCEVertex, Key, [Key])]
            rexpsFor (Module _ _ mn _ _ _ _ _ _) = case M.lookup mn rexp of
              Nothing -> []
              Just ids -> (\i -> (ReExportedVertex (mkQualified i parent)
                                , mkQualified i parent
                                , [mkQualified i mn])
                         ) <$> ids
  
  
        isReExported :: M.Map ModuleName [Ident] -> Module Ann -> Bool
        isReExported rexps (Module _ _ name _ _ _ _ _ _) =
          M.member name rexps
  
        -- | Find dependencies of an expression.
        deps :: Expr Ann -> [Key]
        deps = traverseExpr
          where
            -- | Build graph from qualified identifiers.
            onExpr :: Expr Ann -> [Key]
            onExpr (Var _ i) = [i | isQualified i]
            onExpr _ = []
  
            traverseExpr :: Expr Ann -> [Key]
            traverseExpr v@(Literal _ l) =
              foldl'
                (++)
                (onExpr v)
                (map traverseExpr (extractLiteral l))
            traverseExpr v@Constructor {} = onExpr v
            traverseExpr v@(Accessor _ _ e1) = onExpr v ++ traverseExpr e1
            traverseExpr v@(ObjectUpdate _ obj vs) =
              foldl'
                (++)
                (onExpr v ++ traverseExpr obj)
                (map (traverseExpr . snd) vs)
            traverseExpr v@(Abs _ _ e1) = onExpr v ++ traverseExpr e1
            traverseExpr v@(App _ e1 e2) =
              onExpr v ++ traverseExpr e1 ++ traverseExpr e2
            traverseExpr v@(Var _ _) = onExpr v
            traverseExpr v@(Case _ vs alts) =
              foldl'
                (++)
                (foldl' (++) (onExpr v) (map traverseExpr vs))
                (map onCaseAlternative alts)
            traverseExpr v@(Let _ ds e1) =
              foldl'
                (++)
                (onExpr v)
                (map onBind ds)
              ++
                traverseExpr e1
  
            onBind :: Bind Ann -> [Key]
            onBind (NonRec _ _ e) = traverseExpr e
            onBind (Rec es) = concatMap (traverseExpr . snd) es
  
            onBinder :: Binder Ann -> [Key]
            onBinder b@(LiteralBinder _ l) =
              foldl'
                (++)
                (onBinder' b)
                (map onBinder (extractLiteral l))
            onBinder b@(ConstructorBinder _ _ _ bs) =
              foldl'
                (++)
                (onBinder' b)
                (map onBinder bs)
            onBinder b@(NamedBinder _ _ b1) = onBinder' b ++ onBinder b1
            onBinder b = onBinder' b
  
            onBinder' :: Binder Ann -> [Key]
            onBinder' (ConstructorBinder _ _ c _) =
              [fmap (Ident . runProperName) c]
            onBinder' _ = []
  
            onCaseAlternative :: CaseAlternative Ann -> [Key]
            onCaseAlternative (CaseAlternative bs (Right val)) =
              concatMap onBinder bs ++ traverseExpr val
            onCaseAlternative (CaseAlternative bs (Left gs)) =
              concat
                (map
                   onBinder
                   bs ++ concatMap
                           (\(grd, val) ->
                              [traverseExpr grd, traverseExpr val]) gs)
  
            -- @f@ is either 'Expr' or 'Binder'
            extractLiteral :: Literal (f Ann) -> [f Ann]
            extractLiteral (ArrayLiteral xs) = xs
            extractLiteral (ObjectLiteral xs) = map snd xs
            extractLiteral _ = []
  
  
      -- | Vertices corresponding to the entry points which we want to keep.
      entryPointVertices :: [Vertex]
      entryPointVertices = catMaybes $ do
        (_, k, _) <- verts
        guard $ k `elem` entryPoints
        return (vertexForKey k)
  
      -- | The list of reachable vertices grouped by module name.
      reachableList :: [[(DCEVertex, Key, [Key])]]
      reachableList
        = groupBy (\(_, k1, _) (_, k2, _) -> getQual k1 == getQual k2)
        $ sortBy (\(_, k1, _) (_, k2, _) -> getQual k1 `compare` getQual k2)
        $ map keyForVertex (concatMap (reachable graph) entryPointVertices)
  
      reachableInModule :: [([(DCEVertex, Key, [Key])], Module Ann)]
      reachableInModule = do
        vs <- reachableList
        m <- modules
        guard (getModuleName vs == Just (moduleName m))
        return (vs, m)
  
      getModuleName :: [(DCEVertex, Key, [Key])] -> Maybe ModuleName
      getModuleName [] = Nothing
      getModuleName ((_, k, _) : _) = getQual k
  
  
  -- | Dead code elimination of local identifiers in `Bind`s, which detects and
  -- removes unused bindings.
  --
  runBindDeadCodeElimination :: Bind Ann -> Bind Ann
  runBindDeadCodeElimination = go
    where
    (go, _, _) = everywhereOnValues id exprFn id
  
    exprFn :: Expr Ann -> Expr Ann
    exprFn (Let ann bs ex) =
      let nbs = foldr' bindFn [] bs
      in if null nbs
        then ex
        else Let ann nbs ex
  
      where
      bindFn :: Bind Ann -> [Bind Ann] -> [Bind Ann]
      bindFn b@(NonRec _ i _) r | i `elem` reachableIdents = b : r
                                | otherwise = r
      bindFn (Rec l) r =
        let l' = filter (\((_, i), _) -> i `elem` reachableIdents) l
        in if null l'
            then r
            else Rec l' : r
  
      (graph, keyForVertex, vertexForKey) = graphFromEdges verts
  
      -- | Build list of vertices
      --
      -- Under the assumption that all identifiers are unique, which is
      -- fullfiled by PureScript.
      verts :: [(Ident, Ident, [Ident])]
      verts = do
        let bes = unBind `concatMap` bs
        (i, e) <- bes
        let deps = fst `map` filter fn bes
              where
                fn (i', _) = i' /= i && isUsedInExpr i' e
        return (i, i, deps)
  
      entryPointVertices :: [Vertex]
      entryPointVertices = catMaybes $ do
        (_, i, _) <- verts
        guard $ isUsedInExpr i ex
        return (vertexForKey i)
  
      reachableIdents = map fn $ reachable graph `concatMap` entryPointVertices
        where
          fn v = case keyForVertex v of (_, i, _) -> i
    exprFn e = e
  
    isUsedInExpr :: Ident -> Expr Ann -> Bool
    isUsedInExpr i (Literal _ (ArrayLiteral es))
      = any (isUsedInExpr i) es
    isUsedInExpr i (Literal _ (ObjectLiteral es))
      = any (isUsedInExpr i . snd) es
    isUsedInExpr _ (Literal _ _) = False
    isUsedInExpr i (Constructor _ _ _ is) = i `elem` is
    isUsedInExpr i (Accessor _ _ e) = isUsedInExpr i e
    isUsedInExpr i (ObjectUpdate _ e ups)
      = isUsedInExpr i e || any (isUsedInExpr i . snd) ups
    isUsedInExpr i (App _ (Abs _ i' e) r)
      = if isUsedInExpr i r
          then isUsedInExpr i' e
          else isUsedInExpr i e
    isUsedInExpr i (App _ l r) = isUsedInExpr i l || isUsedInExpr i r
    isUsedInExpr i (Abs _ i' e) = i /= i' && isUsedInExpr i e
    isUsedInExpr i (Var _ qi@(Qualified by _)) = qi == Qualified by i
    isUsedInExpr i (Case _ es alts)
      = any (isUsedInExpr i) es || any (isUsedInCaseAlternative i) alts
    isUsedInExpr i (Let _ bs e) =
      if shadowed
          then used
          else used || isUsedInExpr i e
      where
      -- |
      -- Check if an identifier is used in bindings and the resulting
      -- expression.  A binding might shadow an identifier.  The first Boolean
      -- value denotes if i is used in any bind expression, the second if it was
      -- shadowed.
      (used, shadowed) = foldl' fn (False, False) (concatMap unBind bs)
  
      fn (u, s) (i', e')
        | s || i == i'  = (u, True)
        | otherwise     = (u || isUsedInExpr i e', False)
  
    isUsedInCaseAlternative i (CaseAlternative bs ee) =
      not (any (\b -> case b of
                      VarBinder _ i'  -> i == i'
                      _               -> False) bs
          )
      &&
        (case ee of
            Right e -> isUsedInExpr i e
            Left es
              -> any (uncurry (||) . (isUsedInExpr i *** isUsedInExpr i)) es)
  