Better Dot rendering of closures

brat/Brat/Compiler.hs

@@ -60,8 +60,8 @@ writeDot :: [FilePath] -> String -> String -> IO ()
 writeDot libDirs file out = do
   env <- runExceptT $ loadFilename root libDirs file
   -- Discard captureSets; perhaps we could incorporate into the graph
-  (_, _, _, graph, _) <- eitherIO env
-  writeFile out (toDotString graph)
+  (_, _, _, graph, cs) <- eitherIO env
+  writeFile out (toDotString graph cs)
 {-
  where
   isMain (PrefixName [] "main", _) = True

brat/Brat/Dot.hs

@@ -1,5 +1,6 @@
 module Brat.Dot (toDotString) where
 
+import Brat.Checker.Monad (CaptureSets)
 import Brat.Naming
 import Brat.Graph
 import Brat.Syntax.Common
@@ -11,7 +12,9 @@ import qualified Data.GraphViz.Printing as GV
 import qualified Data.GraphViz.Attributes.Complete as GV
 
 import qualified Data.Map as M
+import qualified Data.Set as S
 import Data.Text.Lazy (pack, unpack)
+import Data.Maybe (fromMaybe)
 import Data.Bifunctor (first)
 import Data.Graph (reachable, transposeG)
 import Data.Maybe (fromJust)
@@ -34,8 +37,8 @@ instance Show EdgeType where
   show (GraphEdge ty) = show ty
 
 
-toDotString :: Graph -> String
-toDotString (ns,ws) = unpack . GV.printDotGraph $ GV.graphElemsToDot params verts edges
+toDotString :: Graph -> CaptureSets -> String
+toDotString (ns,ws) cs = unpack . GV.printDotGraph $ GV.graphElemsToDot params verts edges
  where
   verts :: [(Name', Node)]
   verts = first Name' <$> M.toList ns
@@ -50,26 +53,31 @@ toDotString (ns,ws) = unpack . GV.printDotGraph $ GV.graphElemsToDot params vert
   getRefEdge :: Name' -> Node -> [(Name', Name', EdgeType)]
   getRefEdge x (BratNode (Eval (Ex y _)) _ _) = [(Name' y, x, EvalEdge)]
   getRefEdge x (KernelNode (Splice (Ex y _)) _ _) = [(Name' y, x, EvalEdge)]
-  getRefEdge x (BratNode (Box src _) _ _) = [(x, Name' src, SrcEdge)]
+  getRefEdge x (BratNode (Box src tgt) _ _) = [(x, Name' src, SrcEdge), (x, Name' tgt, SrcEdge)]
   getRefEdge _ _ = []
 
-  -- Map all nodes in a box to the src node
-  clusterMap :: M.Map Name' Name'
+  -- Map from node to cluster. Clusters are identified by their containing Box node.
+  clusterMap :: M.Map Name' Name
   clusterMap = foldr f M.empty verts
    where
     (g, toNode, toVert) = toGraph (ns, ws)
-    f (_, node) m = case node of
-      BratNode (Box src tgt) _ _ ->
-        -- Find all nodes in the box spanned by src and tgt, i.e. all nodes
-        -- reachable from src that can reach tgt
-        let srcReaches = reachable g (fromJust (toVert src))
-            reachesTgt = reachable (transposeG g) (fromJust (toVert tgt))
-            box = Name' . snd3 . toNode <$> (srcReaches ++ reachesTgt) in
-        foldr (`M.insert` Name' src) m box
-      _ -> m
+    f (Name' boxNode, BratNode (Box src tgt) _ _) m =
+      -- Find all nodes in the box spanned by src and tgt, i.e. all nodes
+      -- reachable from src that can reach tgt
+      let srcReaches = reachable g (fromJust (toVert src))
+          reachesTgt = reachable (transposeG g) (fromJust (toVert tgt))
+          nodesUsedInBox = snd3 . toNode <$> (srcReaches ++ reachesTgt)
+          -- exclude nodes that are captured by the box - these are not in the box
+          -- (TODO: we might consider adding extra edges from these to the box itself,
+          --  but for now they'll just have "normal" value edges *entering* the box)
+          captures = fromMaybe M.empty (M.lookup boxNode cs)
+          captureNodes = S.fromList [n | vs <- M.elems captures, (NamedPort (Ex n _) _, _) <- vs]
+          nodesInBox = [Name' n | n <- nodesUsedInBox, S.notMember n captureNodes]
+      in foldr (`M.insert` boxNode) m nodesInBox
+    f _ m = m
 
   -- GV.GraphVisParams vertexType vertexLabelType edgeLabelType clusterType clusterLabelType
-  params :: GV.GraphvizParams Name' Node EdgeType Name' Node
+  params :: GV.GraphvizParams Name' Node EdgeType Name Node
   params = GV.defaultParams {
     GV.fmtNode = \(Name' name, node) -> [
       GV.textLabel (pack $ show name ++ ":\\n" ++ showNodeType node),
@@ -83,13 +91,20 @@ toDotString (ns,ws) = unpack . GV.printDotGraph $ GV.graphElemsToDot params vert
       GV.Style [style label],
       GV.arrowTo (arrow label)
     ],
-    GV.clusterBy = \n@(name, _) -> case clusterMap M.!? name of
-        Just parent -> GV.C parent $ GV.N n
-        Nothing -> GV.N n,
-    GV.clusterID = \(Name' name) -> GV.Str (pack $ show name)
+    GV.clusterBy = \n@(name, _) -> nestClusters name (GV.N n),
+    GV.clusterID = GV.Str . pack . show
   }
 
+  nestClusters :: Name' -> GV.NodeCluster Name (Name', Node) -> GV.NodeCluster Name (Name', Node)
+  nestClusters name n = case clusterMap M.!? name of
+    Nothing -> n
+    -- put n in clust, which may itself be in another cluster
+    Just clust -> nestClusters (Name' clust) (GV.C clust n)
+
   showNodeType :: Node -> String
+  -- Do not repeat the internal links that have been turned into edges
+  showNodeType (BratNode (Box _ _) _ _) = "Box"
+  showNodeType (BratNode (Eval _) _ _) = "Eval"
   showNodeType (BratNode thing _ _) = show thing
   showNodeType (KernelNode thing _ _) = show thing
 

brat/test/compilation/closures.brat

@@ -1,17 +1,2 @@
 f(Nat) -> { Nat -> Nat }
 f(x) = { y => x + y }
-
-g(Nat) -> { Nat -> { Nat -> Nat } }
-g(x) = { y => { z => x + y + z } }
-
-
-h(Nat) -> { Nat -> Nat }
-h(x) = let y = x in { z => x + y + z }
-
-ff(Vec(Nat,2)) -> { Nat -> Nat }
-ff([a,b]) = { y => a + b*y }
-
-ext "to_float" i2f :: {Int -> Float}
-
-fi(Float) -> { Int -> Float }
-fi(x) = { y => x + i2f(y) }