add Stg.Tickish utility module to collect source locations for STG IR points

Author: csabahruska

external-stg/app/stgapp.hs

@@ -44,6 +44,7 @@ modes = subparser
     <> mode "undef" undefMode  (progDesc "print list of undefined foreign symbols")
     <> mode "link" linkMode  (progDesc "link cbits.so for the applications with used foreign functions")
     <> mode "hi-list"  hiListMode  (progDesc "whole program interface file list")
+    <> mode "srcpaths"  srcpathMode  (progDesc "print module source filepaths")
     )
   where
     mode :: String -> Parser a -> InfoMod a -> Mod CommandFields a
@@ -209,5 +210,32 @@ modes = subparser
               True  -> printf "OK: %s\n" modModuleName
               False -> printf "MISSING: %s\n" hiName
 
+    srcpathMode :: Parser (IO ())
+    srcpathMode =
+        run <$> fullpakFile
+      where
+        run fname = do
+          moduleList <- loadModules fname
+          -- print source filepaths
+          forM_ moduleList $ \m -> case moduleSourceFilePath m of
+            Nothing       -> pure ()
+            Just srcPath  -> printf "source filepath: %s %s %s\n"
+                              (BS8.unpack $ getUnitId $ moduleUnitId m)
+                              (BS8.unpack $ getModuleName $ moduleName m)
+                              (BS8.unpack srcPath)
+
+          -- report empty moduleSourceFilePath
+          forM_ [m | m <- moduleList, isNothing $ moduleSourceFilePath m] $ \m -> do
+            printf "missing source filepath for: %s %s\n" (BS8.unpack $ getUnitId $ moduleUnitId m) (BS8.unpack $ getModuleName $ moduleName m)
+
+          -- report ambiguous moduleSourceFilePath
+          let moduleMaps  = [Map.singleton srcPath (1, [m]) | m <- moduleList, srcPath <- maybeToList $ moduleSourceFilePath m]
+              duplicates  = Map.filter (\(n, _) -> n > 1) $ Map.unionsWith (\(n1, l1) (n2, l2) -> (n1 + n2, l1 ++ l2)) moduleMaps
+          forM_ (Map.toList duplicates) $ \(srcPath, (_, mods)) -> forM_ mods $ \m -> do
+            printf "duplicate source filepath: %s %s %s\n"
+              (BS8.unpack $ getUnitId $ moduleUnitId m)
+              (BS8.unpack $ getModuleName $ moduleName m)
+              (BS8.unpack srcPath)
+
 main :: IO ()
 main = join $ execParser $ info (helper <*> modes) mempty

external-stg/external-stg.cabal

@@ -16,6 +16,7 @@ library
   exposed-modules:
     Stg.Pretty
     Stg.IRLocation
+    Stg.Tickish
     Stg.Reconstruct
     Stg.Deconstruct
     Stg.Fullpak

external-stg/lib/Stg/Tickish.hs

@@ -0,0 +1,60 @@
+{-# LANGUAGE LambdaCase #-}
+module Stg.Tickish where
+
+import Control.Monad.RWS hiding (Alt)
+
+import Stg.Syntax
+import Stg.IRLocation
+
+type M = RWS (Maybe StgPoint) [(StgPoint, Tickish)] ()
+
+withStgPoint :: StgPoint -> M () -> M ()
+withStgPoint sp = local (const $ Just sp)
+
+getStgPoint :: M StgPoint
+getStgPoint = ask >>= \case
+  Nothing -> error "missing stg point"
+  Just sp -> pure sp
+
+visitTopBinding :: TopBinding -> M ()
+visitTopBinding = \case
+  StgTopLifted b    -> visitBinding b
+  StgTopStringLit{} -> pure ()
+
+visitBinding :: Binding -> M ()
+visitBinding = \case
+  StgNonRec b r  -> visitRhs b r
+  StgRec bs      -> mapM_ (uncurry visitRhs) bs
+
+visitRhs :: Binder -> Rhs -> M ()
+visitRhs rhsBinder = \case
+  StgRhsClosure _ _ _ e -> withStgPoint (SP_RhsClosureExpr $ Id rhsBinder) $ visitExpr e
+  StgRhsCon{}           -> pure ()
+
+visitExpr :: Expr -> M ()
+visitExpr expr = do
+  stgPoint <- getStgPoint
+  case expr of
+    StgLit{}            -> pure ()
+    StgApp{}            -> pure ()
+    StgOpApp{}          -> pure ()
+    StgConApp{}         -> pure ()
+    StgCase x b _ alts  -> do
+                            withStgPoint (SP_CaseScrutineeExpr $ Id b) $ visitExpr x
+                            sequence_ [visitAlt (Id b) idx a | (idx, a) <- zip [0..] alts]
+    StgLet b e          -> do
+                            visitBinding b
+                            withStgPoint (SP_LetExpr stgPoint) $ visitExpr e
+    StgLetNoEscape b e  -> do
+                            visitBinding b
+                            withStgPoint (SP_LetNoEscapeExpr stgPoint) $ visitExpr e
+    StgTick tickish e   -> do
+                            tell [(stgPoint, tickish)]
+                            visitExpr e
+
+visitAlt :: Id -> Int -> Alt -> M ()
+visitAlt scrutId idx (Alt _con _bndrs rhs) = do
+  withStgPoint (SP_AltExpr scrutId idx) $ visitExpr rhs
+
+collectTickish :: Module -> [(StgPoint, Tickish)]
+collectTickish m = snd $ evalRWS (mapM_ visitTopBinding $ moduleTopBindings m) Nothing ()