Add --external flag to stack dot

src/Stack/Build.hs

@@ -14,7 +14,8 @@
 
 module Stack.Build
   (build
-  ,clean)
+  ,clean
+  ,withLoadPackage)
   where
 
 import           Control.Monad
@@ -101,7 +102,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,185 @@
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE TemplateHaskell #-}
-module Stack.Dot where
-
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE OverloadedStrings #-}
+module Stack.Dot (dot
+                 ,DotOpts(..)
+                 ,dotOptsParser
+                 ,resolveDependencies
+                 ,printGraph
+                 ) where
 
-import           Control.Monad (when)
+import           Control.Monad (void)
 import           Control.Monad.Catch (MonadCatch)
-import           Control.Monad.IO.Class (MonadIO)
+import           Control.Monad.IO.Class
 import           Control.Monad.Logger (MonadLogger, logInfo)
 import           Control.Monad.Reader (MonadReader)
+import           Control.Monad.Trans.Control (MonadBaseControl)
 import qualified Data.Foldable as F
-import           Data.Monoid ((<>))
+import qualified Data.HashSet as HashSet
+import           Data.Map (Map)
 import qualified Data.Map as Map
+import           Data.Monoid ((<>))
+import           Data.Set (Set)
 import qualified Data.Set as Set
-import qualified Data.Text as T
+import           Data.Text (Text)
+import qualified Data.Text as Text
+import qualified Data.Traversable as T
+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.Constants
 import           Stack.Package
 import           Stack.Types
 
+-- | Options record for `stack dot`
+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
+    }
+
+-- | Parser for arguments to `stack dot`
+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)")))
+
+-- | Visualize the project's dependencies as a graphviz graph
+dot :: (HasEnvConfig env
+       ,HasHttpManager env
+       ,MonadBaseControl IO m
+       ,MonadCatch m
+       ,MonadIO m
+       ,MonadLogger m
+       ,MonadReader env m
+       )
+    => DotOpts
+    -> m ()
+dot dotOpts = do
+    (locals,_,_) <- loadLocals defaultBuildOpts Map.empty
+    (_,_,_,sourceMap) <- loadSourceMap defaultBuildOpts
+    let graph = Map.fromList (localDependencies dotOpts locals)
+    menv <- getMinimalEnvOverride
+    resultGraph <- withLoadPackage menv (\loader -> do
+      let depLoader = createDepLoader sourceMap (fmap3 packageAllDeps loader)
+      liftIO $ resolveDependencies (dotDependencyDepth dotOpts) graph depLoader)
+    printGraph dotOpts locals (if dotIncludeBase dotOpts
+                                 then resultGraph
+                                 else filterOutDepsOnBase resultGraph)
+  where filterOutDepsOnBase = Map.filterWithKey (\k _ -> show k /= "base") .
+                              fmap (Set.filter ((/= "base") . show))
+        -- fmap a function over the result of a function with 3 arguments
+        fmap3 :: Functor f => (d -> e) -> (a -> b -> c -> f d) -> (a -> b -> c -> f e)
+        fmap3 f g a b c = f <$> g a b c
+
+-- | Resolve the dependency graph up to (Just depth) or until fixpoint is reached
+resolveDependencies :: (Applicative m, Monad m)
+                    => Maybe Int
+                    -> Map PackageName (Set PackageName)
+                    -> (PackageName -> m (Set PackageName))
+                    -> m (Map PackageName (Set PackageName))
+resolveDependencies (Just 0) graph _ = return graph
+resolveDependencies limit graph loadPackageDeps = 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 <- T.traverse (\name -> (name,) <$> loadPackageDeps name) (F.toList next)
+       resolveDependencies (subtract 1 <$> limit)
+                      (Map.unionWith Set.union graph (Map.fromList x))
+                      loadPackageDeps
+
+-- | Given a SourceMap and a dependency loader, load the set of dependencies for a package
+createDepLoader :: Applicative m
+                => Map PackageName PackageSource
+                -> (PackageName -> Version -> Map FlagName Bool -> m (Set PackageName))
+                -> PackageName
+                -> m (Set PackageName)
+createDepLoader sourceMap loadPackageDeps pkgName =
+  case Map.lookup pkgName sourceMap of
+    Just (PSLocal lp) -> pure (packageAllDeps (lpPackage lp))
+    Just (PSUpstream version _ flags) -> loadPackageDeps pkgName version flags
+    Nothing -> pure Set.empty
+
+-- | Resolve the direct (depth 0) external dependencies of the given local packages
+localDependencies :: DotOpts -> [LocalPackage] -> [(PackageName,Set PackageName)]
+localDependencies dotOpts locals = 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
+
+-- | Print a graphviz graph of the edges in the Map and highlight the given local packages
+printGraph :: (Applicative m, MonadLogger m)
+           => DotOpts
+           -> [LocalPackage]
+           -> Map PackageName (Set PackageName)
+           -> m ()
+printGraph dotOpts locals graph = do
+  $logInfo "strict digraph deps {"
+  printLocalNodes dotOpts locals
+  printLeaves graph
+  void (Map.traverseWithKey printEdges graph)
+  $logInfo "}"
+
+-- | Print the local nodes with a different style depending on options
+printLocalNodes :: (F.Foldable t, MonadLogger m)
+                => DotOpts
+                -> t LocalPackage
+                -> m ()
+printLocalNodes dotOpts locals = $logInfo (Text.intercalate "\n" lpNodes)
+  where applyStyle :: Text -> Text
+        applyStyle n = if dotIncludeExternal dotOpts
+                         then n <> " [style=dashed];"
+                         else n <> " [style=solid];"
+        lpNodes :: [Text]
+        lpNodes = map (applyStyle . nodeName . packageName . lpPackage) (F.toList locals)
+
+-- | Print nodes without dependencies
+printLeaves :: (Applicative m, MonadLogger m) => Map PackageName (Set PackageName) -> m ()
+printLeaves = F.traverse_ printLeaf . Map.keysSet . Map.filter Set.null
+
+-- | `printDedges p ps` prints an edge from p to every ps
+printEdges :: (Applicative m, MonadLogger m) => PackageName -> Set PackageName -> m ()
+printEdges package deps = F.for_ deps (printEdge package)
+
+-- | Print an edge between the two package names
+printEdge :: MonadLogger m => PackageName -> PackageName -> m ()
+printEdge from to = $logInfo (Text.concat [ nodeName from, " -> ", nodeName to, ";"])
+
 -- | 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 "}"
+nodeName :: PackageName -> Text
+nodeName name = "\"" <> Text.pack (packageNameString name) <> "\""
+
+-- | Print a node with no dependencies
+printLeaf :: MonadLogger m => PackageName -> m ()
+printLeaf package = $logInfo . Text.concat $
+  if isWiredIn package
+    then ["{rank=max; ", nodeName package, " [shape=box]; };"]
+    else ["{rank=max; ", nodeName package, "; };"]
+
+-- | Check if the package is wired in (shipped with) ghc
+isWiredIn :: PackageName -> Bool
+isWiredIn = (`HashSet.member` wiredInPackages)

src/main/Main.hs

@@ -165,7 +165,7 @@ main = withInterpreterArgs stackProgName $ \args isInterpreter ->
              addCommand "dot"
                         "Visualize your project's dependency graph using Graphviz dot"
                         dotCmd
-                        (pure ())
+                        dotOptsParser
              addCommand "exec"
                         "Execute a command"
                         execCmd
@@ -892,5 +892,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)

src/test/Stack/DotSpec.hs

@@ -0,0 +1,92 @@
+{-# LANGUAGE OverloadedStrings #-}
+-- | Test suite for Stack.Dot
+module Stack.DotSpec where
+
+import           Data.ByteString.Char8 (ByteString)
+import           Data.Functor.Identity
+import qualified Data.Map as Map
+import           Data.Maybe (fromMaybe)
+import           Data.Set (Set)
+import qualified Data.Set as Set
+import           Options.Applicative (execParserPure,idm,prefs,info,getParseResult)
+import           Stack.Types
+import           Test.Hspec
+
+import           Stack.Dot
+
+spec :: Spec
+spec = do
+  let graph =
+         Map.mapKeys pkgName
+       . fmap (Set.map pkgName)
+       . Map.fromList $ [("one",Set.fromList ["base","free"])
+                        ,("two",Set.fromList ["base","free","mtl","transformers","one"])
+                        ]
+  describe "Stack.Dot" $ do
+    it "does nothing if depth is 0" $
+      resolveDependencies (Just 0) graph stubLoader `shouldBe` return graph
+
+    it "with depth 1, more dependencies are resolved" $ do
+      let graph' = Map.insert (pkgName "cycle") (Set.singleton (pkgName "cycle")) graph
+          resultGraph = runIdentity (resolveDependencies (Just 0) graph stubLoader)
+          resultGraph' = runIdentity (resolveDependencies (Just 1) graph' stubLoader)
+      Map.size resultGraph < Map.size resultGraph' `shouldBe` True
+
+    it "cycles are ignored" $ do
+       let graph' = Map.insert (pkgName "cycle") (Set.singleton (pkgName "cycle")) graph
+           resultGraph = resolveDependencies Nothing graph stubLoader
+           resultGraph' = resolveDependencies Nothing graph' stubLoader
+       fmap Map.size resultGraph' `shouldBe` fmap ((+1) . Map.size) resultGraph
+
+  where graphElem e graph = Set.member e . Set.unions . Map.elems $ graph
+
+{- Helper functions below -}
+
+-- Unsafe internal helper to create a package name
+pkgName :: ByteString -> PackageName
+pkgName = fromMaybe failure . parsePackageName
+  where
+   failure = (error "Internal error during package name creation in DotSpec.pkgName")
+
+-- Stub, simulates the function to load package dependecies
+stubLoader :: PackageName -> Identity (Set PackageName)
+stubLoader name = return . Set.fromList . map pkgName $ case show name of
+  "StateVar" -> ["stm","transformers"]
+  "array" -> []
+  "bifunctors" -> ["semigroupoids","semigroups","tagged"]
+  "binary" -> ["array","bytestring","containers"]
+  "bytestring" -> ["deepseq","ghc-prim","integer-gmp"]
+  "comonad" -> ["containers","contravariant","distributive"
+               ,"semigroups","tagged","transformers","transformers-compat"
+               ]
+  "cont" -> ["StateVar","semigroups","transformers","transformers-compat","void"]
+  "containers" -> ["array","deepseq","ghc-prim"]
+  "deepseq" -> ["array"]
+  "distributive" -> ["ghc-prim","tagged","transformers","transformers-compat"]
+  "free" -> ["bifunctors","comonad","distributive","mtl"
+            ,"prelude-extras","profunctors","semigroupoids"
+            ,"semigroups","template-haskell","transformers"
+            ]
+  "ghc" -> []
+  "hashable" -> ["bytestring","ghc-prim","integer-gmp","text"]
+  "integer" -> []
+  "mtl" -> ["transformers"]
+  "nats" -> []
+  "one" -> ["free"]
+  "prelude" -> []
+  "profunctors" -> ["comonad","distributive","semigroupoids","tagged","transformers"]
+  "semigroupoids" -> ["comonad","containers","contravariant","distributive"
+                     ,"semigroups","transformers","transformers-compat"
+                     ]
+  "semigroups" -> ["bytestring","containers","deepseq","hashable"
+                  ,"nats","text","unordered-containers"
+                  ]
+  "stm" -> ["array"]
+  "tagged" -> ["template-haskell"]
+  "template" -> []
+  "text" -> ["array","binary","bytestring","deepseq","ghc-prim","integer-gmp"]
+  "transformers" -> []
+  "two" -> ["free","mtl","one","transformers"]
+  "unordered" -> ["deepseq","hashable"]
+  "void" -> ["ghc-prim","hashable","semigroups"]
+  _ -> []

stack.cabal

@@ -226,6 +226,8 @@ test-suite stack-test
                 , resourcet
                 , Cabal
                 , text
+                , optparse-applicative
+                , bytestring
   default-language:    Haskell2010
 
 test-suite stack-integration-test