Haskell source explorer aide

explore.hs

{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import           Control.Monad
import           Control.Monad.Writer
import           Data.List
import           Data.Map                        (Map)
import qualified Data.Map                        as M
import           Data.Monoid
import           Language.Haskell.Exts           (parseFileWithMode)
import           Language.Haskell.Exts.Extension
import           Language.Haskell.Exts.Parser
import           Language.Haskell.Exts.Syntax
import           Lens.Family                     (over)
import           Lens.Family.Stock               (_2)
import           Pipes
import           Pipes.Parse
import qualified Pipes.Prelude                   as P
import           System.Directory
import           System.Environment
import           System.Exit
import           System.FilePath
import           System.IO
import           Text.Printf

type Imports = (ModuleName, [ModuleName])

main :: IO ()
main = getArgs >>= \case
    [dir] -> doesDirectoryExist dir >>= \case
        True -> do
           (importCounts, filenames) <- runWriterT (evalStateT drawAll (imports dir))

           mapM_ (display filenames)
            . sortBy (\(m1,n1) (m2,n2) -> compare n1 n2 <> compare m1 m2)
            . counts $ importCounts
        False -> printUsage
    _ -> printUsage
  where
    imports :: FilePath -> Producer Imports (WriterT (Map ModuleName FilePath) IO) ()
    imports path = emitDirectory path
                   >-> P.filter (isSuffixOf "hs")
                   >-> P.mapM hsFileImports
                   >-> P.concat

    printUsage :: IO a
    printUsage = do
        prog <- getProgName
        hPutStrLn stderr (printf "Usage: %s <src-folder>" prog)
        exitFailure

    display :: Map ModuleName FilePath -> (ModuleName, Int) -> IO ()
    display filenames (name, n) = printf "%d\t%s\t(%s)\n" n (unModuleName name) (filenames M.! name)

    unModuleName :: ModuleName -> String
    unModuleName (ModuleName n) = n

-- Given a list of elements and their links to other elements, construct a
-- list containing the number of links each element has, where a link to
-- another element counts as 1 + the number of links that element has, and
-- a link only counts if it's to another element in the input list.
--
-- Makes perfect sense, right? In the context of counting imports, we *do*
-- want to count things like "MyModule.Foo imports MyModule.Bar", but we
-- don't care that "MyModule.Foo imports Data.List". Thus, an input list of
--
-- [("MyModule.Foo", ["MyModule.Bar","Data.List"]), ("MyModule.Bar", ["Data.Monoid"])]
--
-- will return
--
-- [("MyModule.Foo", 1), ("MyModule.Bar", 0)]
--
-- Note that if "Foo" imports "Bar" and "Baz", and "Bar" also imports "Baz",
-- "Baz" will count twice towards "Foo"'s import count. Oh well.
--
-- This crashes at runtime in the presence of loops, i.e. mutually recursive modules =(
counts :: forall a. Ord a => [(a, [a])] -> [(a, Int)]
counts ys = M.toList countMap
  where
    count :: (a, [a]) -> (a, Int)
    count = over _2 (foldl' step 0 . map (`M.lookup` countMap))
      where
        step :: Int -> Maybe Int -> Int
        step n = maybe n (+(n+1))

    countMap :: Map a Int
    countMap = M.fromList (map count ys)

-- | Emit all non-hidden files in a directory tree.
emitDirectory :: MonadIO m => FilePath -> Producer FilePath m ()
emitDirectory dir =
    liftIO (filterOutDotfiles <$> getDirectoryContents dir) >>= mapM_ go
  where
    go :: MonadIO m => FilePath -> Producer FilePath m ()
    go path = do
        let full_path = dir </> path
        liftIO (doesDirectoryExist full_path) >>= \case
            True  -> emitDirectory full_path
            False -> yield full_path

    filterOutDotfiles :: [FilePath] -> [FilePath]
    filterOutDotfiles = filter (not . isPrefixOf ".")

-- | Get all modules imported by the given file (assumed to be a haskell module).
-- Returns Nothing if the file failed to parse.
hsFileImports :: (MonadIO m, MonadWriter (Map ModuleName FilePath) m) => FilePath -> m (Maybe Imports)
hsFileImports path = do
    let parseMode = defaultParseMode
                        { parseFilename = path
                        , extensions    = glasgowExts
                        , fixities      = Nothing
                        }

    liftIO (parseFileWithMode parseMode path) >>= \case
        ParseOk (Module _ name _ _ _ imports _) -> do
            tell (M.singleton name path)
            return (Just (name, map importModule imports))
        ParseFailed loc err -> do
            liftIO $ hPutStrLn stderr (printf "Failed to parse %s: %s (%s)" path err (show loc))
            return Nothing