Improve output of juvix dev import-tree print

src/Juvix/Compiler/Concrete/Print/Base.hs

@@ -1257,7 +1257,7 @@ instance PrettyPrint ImportTreeStats where
     header "Import Tree Statistics:"
     header "======================="
     itemize
-      [ noLoc "Total number of modules:" <+> noLoc (pretty _importTreeStatsTotalModules),
+      [ noLoc "Total number of nodes:" <+> noLoc (pretty _importTreeStatsTotalModules),
         noLoc "Total number of edges:" <+> noLoc (pretty _importTreeStatsTotalEdges),
         noLoc "Height (longest chain of imports):" <+> noLoc (pretty _importTreeStatsHeight)
       ]
@@ -1268,33 +1268,41 @@ instance PrettyPrint ImportTree where
     header "Import Tree:"
     header "============"
     hardline
-    forM_ (Map.toList importsTable) $ \(pkgRoot, tbl) -> do
+
+    forM_ (Map.toList importsTable) $ \(pkgRoot, tbl :: Map (Path Rel File) (Set ImportNode)) -> do
       annotated AnnImportant (noLoc ("* Package at " <> pretty pkgRoot))
       hardline
+      let pkgNodes :: HashSet ImportNode = fromJust (nodesByRoot ^. at pkgRoot)
+      header ("Nodes (" <> show (length pkgNodes) <> ")")
+      forM_ pkgNodes $ \node -> do
+        noLoc (pMod (node ^. importNodeFile))
+        hardline
+      hardline
+      let numEdges = sum (map length (toList tbl))
+      header ("Edges (" <> show numEdges <> ")")
       forM_ (Map.toList tbl) $ \(fromFile, toFiles) -> do
-        forM_ toFiles $ \toFile -> do
-          let pMod :: Path x File -> Doc Ann
-              pMod = annotate (AnnKind KNameTopModule) . pretty
-              fromMod = pMod fromFile
-              toMod
+        noLoc (pMod fromFile P.<+> annotate AnnKeyword "imports" P.<+> "(" <> pretty (length toFiles) <> "):")
+        hardline
+        indent . itemize . (`map` (toList toFiles)) $ \toFile -> do
+          let toMod
                 | pkgRoot == toFile ^. importNodePackageRoot = pMod (toFile ^. importNodeFile)
                 | otherwise = pMod (toFile ^. importNodeAbsFile)
-          noLoc (fromMod P.<+> annotate AnnKeyword "imports" P.<+> toMod)
-          hardline
+          noLoc toMod
+        hardline
+        unless (null toFiles) hardline
       hardline
     where
+      pMod :: Path x File -> Doc Ann
+      pMod = annotate (AnnKind KNameTopModule) . pretty
+
       allNodes :: [ImportNode]
       allNodes = HashMap.keys (tree ^. importTree)
 
       allRoots :: [Path Abs Dir]
       allRoots = nubSort (map (^. importNodePackageRoot) allNodes)
 
       nodesByRoot :: HashMap (Path Abs Dir) (HashSet ImportNode)
-      nodesByRoot =
-        foldl'
-          (HashMap.unionWith (<>))
-          mempty
-          [hashMap [(node ^. importNodePackageRoot, hashSet [node])] | node <- allNodes]
+      nodesByRoot = importTreeNodesByPackage tree
 
       -- fromPackageRoot -> fromFile -> tofile
       importsTable :: Map (Path Abs Dir) (Map (Path Rel File) (Set ImportNode))

src/Juvix/Compiler/Pipeline/Loader/PathResolver/ImportTree/Base.hs

@@ -6,6 +6,7 @@ module Juvix.Compiler.Pipeline.Loader.PathResolver.ImportTree.Base
     importTreeEdges,
     importTreeNodes,
     importTreeProjectNodes,
+    importTreeNodesByPackage,
     ImportTreeBuilder,
     runImportTreeBuilder,
     ignoreImportTreeBuilder,
@@ -89,6 +90,12 @@ withImportNode fromNode m = do
   (`interpret` m) $ \case
     ImportTreeAddEdge importScan toNode -> internalRegisterEdge importScan fromNode toNode
 
+importTreeNodesByPackage :: ImportTree -> HashMap (Path Abs Dir) (HashSet ImportNode)
+importTreeNodesByPackage tree = run . execState mempty $
+  forM_ (tree ^. importTreeNodes) $ \node ->
+    modify @(HashMap (Path Abs Dir) (HashSet ImportNode))
+      (over (at (node ^. importNodePackageRoot)) (Just . maybe (HashSet.singleton node) (HashSet.insert node)))
+
 importTree :: SimpleGetter ImportTree (HashMap ImportNode (HashSet ImportNode))
 importTree = fimportTree