ghc-mod/Language/Haskell/GhcMod/HomeModuleGraph.hs

-- ghc-mod: Making Haskell development *more* fun
-- Copyright (C) 2015  Daniel Gröber <dxld ÄT darkboxed DOT org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE ScopedTypeVariables, RecordWildCards #-}
module Language.Haskell.GhcMod.HomeModuleGraph (
   GmModuleGraph(..)
 , ModulePath(..)
 , mkFileMap
 , mkModuleMap
 , mkMainModulePath
 , findModulePath
 , findModulePathSet
 , fileModuleName
 , homeModuleGraph
 , updateHomeModuleGraph
 , reachable
 , moduleGraphToDot
 ) where

import Bag
import DriverPipeline hiding (unP)
import ErrUtils
import Exception
import FastString
import Finder
import GHC
import HscTypes
import Lexer
import MonadUtils hiding (foldrM)
import Parser
import SrcLoc
import StringBuffer

import Control.Arrow ((&&&))
import Control.Monad
import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT)
import Control.Monad.State.Strict (execStateT)
import Control.Monad.State.Class
import Data.Maybe
import Data.Monoid
import Data.Map.Strict  (Map)
import qualified Data.Map.Strict  as Map
import Data.Set (Set)
import qualified Data.Set as Set
import System.FilePath

import Language.Haskell.GhcMod.Logging
import Language.Haskell.GhcMod.Logger
import Language.Haskell.GhcMod.Monad.Types
import Language.Haskell.GhcMod.Types

-- | Turn module graph into a graphviz dot file
--
-- @dot -Tpng -o modules.png modules.dot@
moduleGraphToDot :: GmModuleGraph -> String
moduleGraphToDot GmModuleGraph { gmgGraph } =
    "digraph {\n" ++ concatMap edges (Map.toList graph) ++ "}\n"
 where
   graph = Map.map (Set.mapMonotonic mpPath)
         $ Map.mapKeysMonotonic mpPath gmgGraph
   edges :: (FilePath, (Set FilePath)) -> String
   edges (f, sf) =
       concatMap (\f' -> "    \""++ f ++"\" -> \""++ f' ++"\"\n") (Set.toList sf)

data S = S {
      sErrors   :: [(ModulePath, ErrorMessages)],
      sWarnings :: [(ModulePath, WarningMessages)],
      sGraph    :: GmModuleGraph
}

defaultS :: S
defaultS = S [] [] mempty

putErr :: MonadState S m
       => (ModulePath, ErrorMessages) -> m ()
putErr e = do
  s <- get
  put s { sErrors = e:sErrors s}

putWarn :: MonadState S m
       => (ModulePath, ErrorMessages) -> m ()
putWarn w = do
  s <- get
  put s { sWarnings = w:sWarnings s}

gmgLookupMP :: MonadState S m => ModulePath -> m (Maybe (Set ModulePath))
gmgLookupMP k = (Map.lookup k . gmgGraph . sGraph) `liftM` get

graphUnion :: MonadState S m => GmModuleGraph -> m ()
graphUnion gmg = do
  s <- get
  put s { sGraph = sGraph s `mappend` gmg }

reachable :: Set ModulePath -> GmModuleGraph -> Set ModulePath
reachable smp0 GmModuleGraph {..} = go smp0
 where
   go smp = let
       δsmp = Set.unions $
                collapseMaybeSet . flip Map.lookup gmgGraph <$> Set.toList smp
       smp' = smp `Set.union` δsmp
    in if smp == smp' then smp' else go smp'

pruneUnreachable :: Set ModulePath -> GmModuleGraph -> GmModuleGraph
pruneUnreachable smp0 gmg@GmModuleGraph {..} = let
    r = reachable smp0 gmg
    rfn = Set.map mpPath r
    rmn = Set.map mpModule r
  in
    GmModuleGraph {
      gmgFileMap = Map.filterWithKey (\k _ -> k `Set.member` rfn) gmgFileMap,
      gmgModuleMap = Map.filterWithKey (\k _ -> k `Set.member` rmn) gmgModuleMap,
      gmgGraph = Map.filterWithKey (\k _ -> k `Set.member` r) gmgGraph
    }

collapseMaybeSet :: Maybe (Set a) -> Set a
collapseMaybeSet = maybe Set.empty id

homeModuleGraph :: (IOish m, GmLog m, GmEnv m)
    => HscEnv -> Set ModulePath -> m GmModuleGraph
homeModuleGraph env smp = updateHomeModuleGraph env mempty smp smp

mkMainModulePath :: FilePath -> ModulePath
mkMainModulePath = ModulePath (mkModuleName "Main")

findModulePath :: HscEnv -> ModuleName -> IO (Maybe ModulePath)
findModulePath env mn = do
    fmap (ModulePath mn) <$> find env mn

findModulePathSet :: HscEnv -> [ModuleName] -> IO (Set ModulePath)
findModulePathSet env mns = do
    Set.fromList . catMaybes <$> findModulePath env `mapM` mns

find :: MonadIO m => HscEnv -> ModuleName -> m (Maybe FilePath)
find env mn = liftIO $ do
  res <- findHomeModule env mn
  case res of
   -- TODO: handle SOURCE imports (hs-boot stuff): addBootSuffixLocn loc
    Found loc@ModLocation { ml_hs_file = Just _ } _mod -> do
        return $ normalise <$> ml_hs_file loc
    _ -> return Nothing

updateHomeModuleGraph :: (IOish m, GmLog m, GmEnv m)
                      => HscEnv
                      -> GmModuleGraph
                      -> Set ModulePath -- ^ Initial set of modules
                      -> Set ModulePath -- ^ Updated set of modules
                      -> m GmModuleGraph
updateHomeModuleGraph env GmModuleGraph {..} smp usmp = do
    -- TODO: It would be good if we could retain information about modules that
    -- stop to compile after we've already successfully parsed them at some
    -- point. Figure out a way to delete the modules about to be updated only
    -- after we're sure they won't fail to parse .. or something. Should probably
    -- push this whole prune logic deep into updateHomeModuleGraph'
   (pruneUnreachable smp . sGraph) `liftM` runS (updateHomeModuleGraph' env usmp)
 where
   runS = flip execStateT defaultS { sGraph = graph' }
   graph' = GmModuleGraph {
       gmgFileMap = Set.foldr (Map.delete . mpPath) gmgFileMap usmp,
       gmgModuleMap = Set.foldr (Map.delete . mpModule) gmgModuleMap usmp,
       gmgGraph = Set.foldr Map.delete gmgGraph usmp
    }

mkFileMap :: Set ModulePath -> Map FilePath ModulePath
mkFileMap smp = Map.fromList $ map (mpPath &&& id) $ Set.toList smp

mkModuleMap :: Set ModulePath -> Map ModuleName ModulePath
mkModuleMap smp = Map.fromList $ map (mpModule &&& id) $ Set.toList smp

updateHomeModuleGraph'
    :: forall m. (MonadState S m, IOish m, GmLog m, GmEnv m)
    => HscEnv
    -> Set ModulePath     -- ^ Initial set of modules
    -> m ()
updateHomeModuleGraph' env smp0 = do
    go `mapM_` Set.toList smp0

 where
   go :: ModulePath -> m ()
   go mp = do
     msmp <- gmgLookupMP mp
     case msmp of
       Just _ -> return ()
       Nothing -> do
           smp <- collapseMaybeSet `liftM` step mp

           graphUnion GmModuleGraph {
               gmgFileMap = mkFileMap smp,
               gmgModuleMap = mkModuleMap smp,
               gmgGraph = Map.singleton mp smp
            }

           mapM_ go (Set.toList smp)

   step :: ModulePath -> m (Maybe (Set ModulePath))
   step mp = runMaybeT $ do
       (dflags, ppsrc_fn) <- MaybeT preprocess'
       src <- liftIO $ readFile ppsrc_fn
       imports mp src dflags
    where
      preprocess' :: m (Maybe (DynFlags, FilePath))
      preprocess' = do
        let fn = mpPath mp
        ep <- liftIO $ withLogger' env $ \setDf -> let
              env' = env { hsc_dflags = setDf (hsc_dflags env) }
           in preprocess env' (fn, Nothing)
        case ep of
          Right (_, x) -> return $ Just x
          Left errs -> do
            -- TODO: Remember these and present them as proper errors if this is
            -- the file the user is looking at.
            gmLog GmWarning "preprocess'" $ vcat $ map strDoc errs
            return Nothing

   imports :: ModulePath -> String -> DynFlags -> MaybeT m (Set ModulePath)
   imports mp@ModulePath {..} src dflags =
       case parseModuleHeader src dflags mpPath of
         Left err -> do
           putErr (mp, err)
           mzero

         Right (ws, lmdl) -> do
           putWarn (mp, ws)
           let HsModule {..} = unLoc lmdl
               mns = map (unLoc . ideclName)
                   $ filter (isNothing . ideclPkgQual)
                   $ map unLoc hsmodImports
           liftIO $ Set.fromList . catMaybes <$> mapM (findModulePath env) mns

fileModuleName :: HscEnv
               -> FilePath
               -> IO (Either ErrorMessages (Maybe ModuleName))
fileModuleName env fn = handle (\(_ :: SomeException) -> return $ Right Nothing) $ do
    src <- readFile fn
    case parseModuleHeader src (hsc_dflags env) fn of
      Left errs -> return (Left errs)
      Right (_, lmdl) -> do
        let HsModule {..} = unLoc lmdl
        return $ Right $ unLoc <$> hsmodName

parseModuleHeader
    :: String         -- ^ Haskell module source text (full Unicode is supported)
    -> DynFlags
    -> FilePath       -- ^ the filename (for source locations)
    -> Either ErrorMessages (WarningMessages, Located (HsModule RdrName))
parseModuleHeader str dflags filename =
   let
       loc  = mkRealSrcLoc (mkFastString filename) 1 1
       buf  = stringToStringBuffer str
   in
   case unP Parser.parseHeader (mkPState dflags buf loc) of

     PFailed sp err   ->
         Left (unitBag (mkPlainErrMsg dflags sp err))

     POk pst rdr_module ->
         let (warns,_) = getMessages pst in
         Right (warns, rdr_module)