Add external dependency visualization to stack dot

- new flag --[no-]external to include external dependencies
- new flag --[no-]include-base to toggle edges to base package
- new flag --depth to limit depth of external dependencies shown

src/Stack/Build.hs

@@ -14,7 +14,8 @@
 
 module Stack.Build
   (build
-  ,clean)
+  ,clean
+  ,withLoadPackage)
   where
 
 import           Control.Monad
@@ -87,7 +88,13 @@ mkBaseConfigOpts bopts = do
         }
 
 -- | Provide a function for loading package information from the package index
-withLoadPackage :: M env m
+withLoadPackage :: ( MonadIO m
+                   , HasHttpManager env
+                   , MonadReader env m
+                   , MonadBaseControl IO m
+                   , MonadCatch m
+                   , MonadLogger m
+                   , HasEnvConfig env)
                 => EnvOverride
                 -> ((PackageName -> Version -> Map FlagName Bool -> IO Package) -> m a)
                 -> m a

src/Stack/Dot.hs

@@ -1,53 +1,193 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE TemplateHaskell #-}
-module Stack.Dot where
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TupleSections #-}
+module Stack.Dot (dot, DotOpts(..), dotOptsParser) where
 
-
-import           Control.Monad (when)
-import           Control.Monad.Catch (MonadCatch)
-import           Control.Monad.IO.Class (MonadIO)
+import           Control.Monad (when,void)
+import           Control.Monad.Catch (MonadCatch,MonadMask)
+import           Control.Monad.IO.Class
 import           Control.Monad.Logger (MonadLogger, logInfo)
 import           Control.Monad.Reader (MonadReader)
-import qualified Data.Foldable as F
-import           Data.Monoid ((<>))
+import           Control.Monad.Trans.Control (MonadBaseControl)
+import           Data.Foldable (for_,toList,Foldable)
+import           Data.Map (Map)
 import qualified Data.Map as Map
+import           Data.Monoid ((<>))
+import           Data.Set (Set)
 import qualified Data.Set as Set
+import           Data.Text (Text)
 import qualified Data.Text as T
+import           Data.Traversable (traverse)
+import           Network.HTTP.Client.Conduit (HasHttpManager)
+import           Options.Applicative
+import           Options.Applicative.Builder.Extra (boolFlags)
+import           Stack.Build (withLoadPackage)
 import           Stack.Build.Source
 import           Stack.Build.Types
 import           Stack.Package
 import           Stack.Types
 
+data DotOpts = DotOpts
+    { dotIncludeExternal :: Bool
+    -- ^ Include external dependencies
+    , dotIncludeBase :: Bool
+    -- ^ Include dependencies on base
+    , dotDependencyDepth :: Maybe Int
+    -- ^ Limit the depth of dependency resolution to (Just n) or continue until fixpoint
+    }
+
+dotOptsParser :: Parser DotOpts
+dotOptsParser = DotOpts <$> includeExternal <*> includeBase <*> depthLimit
+  where includeExternal = boolFlags False
+                                    "external"
+                                    "inclusion of external dependencies"
+                                    idm
+        includeBase = boolFlags True
+                                "include-base"
+                                "inclusion of dependencies on base"
+                                idm
+        depthLimit =
+            optional (option auto
+                             (long "depth" <>
+                              metavar "DEPTH" <>
+                              help ("Limit the depth of dependency resolution " <>
+                                    "(Default: No limit)")))
+
 -- | Convert a package name to a graph node name.
 nodeName :: PackageName -> T.Text
 nodeName name = "\"" <> T.pack (packageNameString name) <> "\""
 
-dot :: (MonadReader env m, HasBuildConfig env, MonadIO m, MonadLogger m, MonadCatch m,HasEnvConfig env)
-    => m ()
-dot = do
-    (locals, _names, _idents) <- loadLocals
-        defaultBuildOpts
-        Map.empty
-    let localNames = Set.fromList $ map (packageName . lpPackage) locals
-
-    $logInfo "digraph deps {"
-    $logInfo "splines=polyline;"
-
-    F.forM_ locals $ \lp -> do
-        let deps = Set.intersection localNames $ packageAllDeps $ lpPackage lp
-        F.forM_ deps $ \dep ->
-            $logInfo $ T.concat
-                [ nodeName $ packageName $ lpPackage lp
-                , " -> "
-                , nodeName dep
-                , ";"
-                ]
-        when (Set.null deps) $
-            $logInfo $ T.concat
-                [ "{rank=max; "
-                , nodeName $ packageName $ lpPackage lp
-                , "}"
-                ]
-
-    $logInfo "}"
+dot :: ( HasEnvConfig env
+       , HasHttpManager env
+       , MonadBaseControl IO m
+       , MonadCatch m
+       , MonadIO m
+       , MonadLogger m
+       , MonadReader env m
+       , MonadMask m
+       )
+    => DotOpts
+    -> m ()
+dot dotOpts = do
+    (locals,_,_) <- loadLocals bOpts Map.empty
+    (_,_,_,sourceMap) <- loadSourceMap bOpts
+    let graph = buildGraph dotOpts locals
+    menv <- getMinimalEnvOverride
+    resultGraph <-
+      withLoadPackage menv (dependencyHull (dotDependencyDepth dotOpts) sourceMap graph)
+    printGraph dotOpts locals (if dotIncludeBase dotOpts
+                                 then resultGraph
+                                 else filterOutDepsOnBase resultGraph)
+  where bOpts = BuildOpts { boptsTargets = []
+                          , boptsLibProfile = False
+                          , boptsExeProfile = False
+                          , boptsEnableOptimizations = Nothing
+                          , boptsHaddock = False
+                          , boptsHaddockDeps = Nothing
+                          , boptsFinalAction = DoNothing
+                          , boptsDryrun = False
+                          , boptsGhcOptions = []
+                          , boptsFlags = Map.empty
+                          , boptsInstallExes = False
+                          , boptsPreFetch = False
+                          , boptsTestArgs = []
+                          , boptsOnlySnapshot = False
+                          , boptsCoverage = False
+                          }
+        filterOutDepsOnBase = Map.filterWithKey (\k _ -> show k /= "base") .
+                              fmap (Set.filter ((/= "base") . show))
+
+dependencyHull :: (MonadIO m
+                  ,HasHttpManager env
+                  ,MonadReader env m
+                  ,MonadBaseControl IO m
+                  ,MonadCatch m
+                  ,MonadLogger m
+                  ,HasEnvConfig env
+                  )
+               => Maybe Int -- ^ Limit the steps to (Just n) or continue until fixpoint
+               -> Map PackageName PackageSource
+               -> Map PackageName (Set PackageName)
+               -> (PackageName -> Version -> Map FlagName Bool -> IO Package)
+               -> m (Map PackageName (Set PackageName))
+dependencyHull (Just 0) _ graph _ = return graph
+dependencyHull limit sourceMap graph loadPackage = do
+  let values = Set.unions (Map.elems graph)
+      keys = Map.keysSet graph
+      next = Set.difference values keys
+  if Set.null next
+     then return graph
+     else do
+       x <- traverse (\name ->
+                          (name,) <$> loadDeps sourceMap name loadPackage) (toList next)
+       dependencyHull (subtract 1 <$> limit)
+                      sourceMap
+                      (Map.unionWith Set.union graph (Map.fromList x))
+                      loadPackage
+
+loadDeps :: ( MonadIO m
+            , HasHttpManager env
+            , MonadReader env m
+            , MonadBaseControl IO m
+            , MonadCatch m
+            , MonadLogger m
+            , HasEnvConfig env
+            )
+         => Map PackageName PackageSource
+         -> PackageName
+         -> (PackageName -> Version -> Map FlagName Bool -> IO Package)
+         -> m (Set PackageName)
+loadDeps sourceMap packageName loadPackage =
+  case Map.lookup packageName sourceMap of
+    Just (PSLocal lp) -> return (packageAllDeps (lpPackage lp))
+    Just (PSUpstream version _ flags) -> do
+      pkg <- liftIO (loadPackage packageName version flags)
+      return (packageAllDeps pkg)
+    Nothing -> return Set.empty
+
+buildGraph :: DotOpts -> [LocalPackage] -> Map PackageName (Set PackageName)
+buildGraph dotOpts locals = Map.fromList edges
+  where
+    edges :: [(PackageName, Set PackageName)]
+    edges = map (\lp -> (packageName (lpPackage lp), deps lp)) locals
+      where deps lp =
+              if dotIncludeExternal dotOpts
+                then Set.delete (lpName lp) (packageAllDeps (lpPackage lp))
+                else Set.intersection localNames (packageAllDeps (lpPackage lp))
+    lpName lp = packageName (lpPackage lp)
+    localNames = Set.fromList $ map (packageName . lpPackage) locals
+
+printGraph :: MonadLogger m
+           => DotOpts
+           -> [LocalPackage]
+           -> Map PackageName (Set PackageName)
+           -> m ()
+printGraph dotOpts locals edges = do
+  $logInfo "strict digraph deps {"
+  printNodes dotOpts locals
+  void (Map.traverseWithKey printEdges edges)
+  $logInfo "}"
+
+printNodes :: (Foldable t, MonadLogger m)
+           => DotOpts
+           -> t LocalPackage
+           -> m ()
+printNodes dotOpts locals = $logInfo (T.intercalate "\n" lpNodes)
+  where style :: Text
+        style = if dotIncludeExternal dotOpts
+                  then " [style=dashed];"
+                  else " [style=solid];"
+        lpNodes :: [Text]
+        lpNodes = map ((<> style) . nodeName . packageName . lpPackage) (toList locals)
+
+
+printEdges :: MonadLogger m => PackageName -> Set PackageName -> m ()
+printEdges package deps = do
+    for_ deps (\dep -> $logInfo (T.concat [ nodeName package, " -> ", nodeName dep, ";"]))
+    when (Set.null deps)
+         ($logInfo (T.concat [ "{rank=max; "
+                             , nodeName package
+                             , "}"
+                             ]))

src/main/Main.hs

@@ -162,7 +162,7 @@ main =
              addCommand "dot"
                         "Visualize your project's dependency graph using Graphviz dot"
                         dotCmd
-                        (pure ())
+                        dotOptsParser
              addCommand "exec"
                         "Execute a command"
                         execCmd
@@ -864,5 +864,5 @@ solverOptsParser = boolFlags False
     idm
 
 -- | Visualize dependencies
-dotCmd :: () -> GlobalOpts -> IO ()
-dotCmd () go = withBuildConfig go ThrowException dot
+dotCmd :: DotOpts -> GlobalOpts -> IO ()
+dotCmd dotOpts go = withBuildConfig go ThrowException (dot dotOpts)