Print out span of call graph

package.yaml

@@ -46,6 +46,7 @@ default-extensions:
   - GADTs
   - GeneralizedNewtypeDeriving
   - LambdaCase
+  - MultiWayIf
   - NamedFieldPuns
   - NamedWildCards
   - NoImplicitPrelude

src/Rubyhs.hs

@@ -10,11 +10,13 @@ import Data.Language.Ruby (Block)
 import Frelude
 import System.Process (readProcess)
 import qualified Data.ByteString.Lazy.Char8 as ByteString
-import Rubyhs.References (references)
+import qualified Rubyhs.References
 import Options.Applicative (Parser)
 import qualified Options.Applicative as Options
 import System.IO.Temp (withSystemTempFile)
 import System.IO (hFlush)
+import Data.Tree (Tree(Node), Forest)
+import qualified Data.Aeson as Aeson
 
 main :: IO ()
 main = do
@@ -22,7 +24,14 @@ main = do
   blocks <- decodeInput @Block targets
   if printAST
   then traverse_ @[] putEncoded blocks
-  else traverse_ (putEncoded . references) blocks
+  else do
+    traverse_ (putEncoded . Rubyhs.References.references) blocks
+    traverse_ (ByteString.putStrLn . encode . toJSONForest . Rubyhs.References.graph) blocks
+
+toJSONForest :: Forest Rubyhs.References.Node -> Aeson.Value
+toJSONForest = Aeson.Object . fromList . fmap go
+  where
+  go (Node x xs) = (Rubyhs.References.prettyContext x, toJSONForest xs)
 
 -- | Decode from files and mappends the stuff from stdin.
 decodeInput :: FromJSON a => [FilePath] -> IO [a]

src/Rubyhs/References.hs

@@ -1,5 +1,5 @@
 {-# LANGUAGE DuplicateRecordFields, OverloadedLists, InstanceSigs #-}
-{-# OPTIONS_GHC -Wall #-}
+-- {-# OPTIONS_GHC -Wall #-}
 module Rubyhs.References
   ( References(entries)
   , Entry(..)
@@ -9,6 +9,9 @@ module Rubyhs.References
   , Result(..)
   , Namespace(..)
   , Context(..)
+  , graph
+  , prettyContext
+  , Node(..)
   ) where
 
 import Frelude
@@ -18,13 +21,25 @@ import Data.Map (Map)
 import qualified Data.Map as Map
 import Data.Kind
 import Control.Monad.State
-import Data.Aeson ((.=))
 import qualified Data.Aeson as Aeson
-import Data.HashMap.Strict (HashMap)
 import Data.Set (Set)
 import qualified Data.Set as Set
 import qualified Data.Text as Text
 import Data.Coerce
+import Data.Graph (Graph, Vertex, Forest)
+import qualified Data.Graph as Graph
+
+span :: G node key -> Forest node
+span (g, f, _) = fmap ((\(x, _, _) -> x) . f) <$> Graph.dff g
+
+type G node key = (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)
+
+graph :: Block -> Forest Node
+graph b = span $ Graph.graphFromEdges $ go <$> toList ys
+  where
+  go :: (Node, Set Node) -> (Node, Text, [Text])
+  go (x, xs) = (x, prettyContext x, prettyContext <$> toList xs)
+  Result ys = references b
 
 data Entry a = Entry
   { node :: a
@@ -99,21 +114,12 @@ data Env = Env
   , context      :: Context
   }
 
-data Result = Result
-  { declarations :: Set Node
-  , applications :: Map Node (Set Node)
-  }
+newtype Result = Result (Map Node (Set Node))
 instance ToJSON Result where
-  toJSON Result{declarations,applications} = Aeson.object
-    [ "declarations" .= declarations
-    , "applications" .= f
-    ]
+  toJSON (Result a) = Aeson.Object $ fromList $ go <$> toList a
     where
-    f :: HashMap Text (Set Node)
-    f = fromList @(HashMap _ _) $ go <$> toList applications
-    -- go :: (Node, Set Node) -> (Text, Set Node)
-    go :: (Node, Set Node) -> (Text, Set Node)
-    go (x, y) = (prettyContext x, y)
+    go :: (Node, Set Node) -> (Text, Aeson.Value)
+    go (x, y) = (prettyContext x, Aeson.toJSON y)
 
 prettyContext :: Node -> Text
 prettyContext = \case
@@ -164,10 +170,11 @@ class References a where
   entries :: MyMonad m => a -> m ()
 
 references :: Block -> Result
-references q = Result{declarations, applications}
+references q = Result $ Map.unionWith mappend applications declarations'
   where
   Env{declarations,applications} = execState (entries @_ @(State Env) q) emptyEnv
   emptyEnv = Env mempty mempty (Context $ NodeModule mempty)
+  declarations' = fromList $ (\x -> (x, mempty)) <$> toList declarations
 
 instance References Block where
   entries :: forall m . MyMonad m => Block -> m ()
@@ -193,7 +200,9 @@ instance References Ruby.Namespace where
           NodeModule namespace -> namespace
         -- TODO Hacky:
         ns = Namespace ((\(Name (Aeson.String t)) -> t) <$> xs)
-      application c (NodeModule $ ctxt `onTop` ns)
+      -- TODO: Broken
+      -- application c (NodeModule $ ctxt `onTop` ns)
+      application c (NodeModule $ ns)
 
 instance References Module where
   entries Module{name, block} = do
@@ -237,7 +246,9 @@ instance References Send where
       ctxt = case c of
         NodeFunction FQN{namespace = ns} -> ns
         NodeModule ns -> ns
-    application c $ NodeFunction $ qual (ctxt `onTop` fromNS namespace) name
+    -- TODO: Broken
+    -- application c $ NodeFunction $ qual (ctxt `onTop` fromNS namespace) name
+    application c $ NodeFunction $ qual (fromNS namespace) name
     where
     fromNS :: Ruby.Namespace -> Namespace
     fromNS (Ruby.Namespace l) = Namespace $ go <$> l