TraverseSources.hs

module Main (main) where

import System.FilePath.Glob (glob)
import Control.Monad (foldM)
import Distribution.Simple.Utils (fromUTF8BS)

import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map

import GHC.Hs.Expr
import GHC.Hs.Extension
import SrcLoc
import RdrName
import Language.Haskell.HLint (parseModuleEx, defaultParseFlags, ModuleEx (..))
import OccName

import qualified Data.Generics as SYB

main :: IO ()
main = do
    packages' <- lines <$> readFile "_data/packages.txt"
    let packages = take 10000 packages'

    putStrLn $ show (length packages) ++ " packages"
    
    res <- foldM withPackage Map.empty packages

    let maxi = maximum res

    Map.foldlWithKey (\acc c n -> acc *> display maxi c n) (pure ()) res

display :: Int -> Char -> Int -> IO ()
display maxi c n = putStrLn $ concat
    [ "  " 
    , [c]
    , " "
    , replicate (l - length (show n)) ' '
    , show n
    , " "
    , replicate (n * 120 `div` maxi) '#'
    ]
  where
    l = length (show maxi)

withPackage :: Map.Map Char Int -> String -> IO (Map.Map Char Int)
withPackage acc pkg = do
    files <- glob (pkg ++ "*/**/*.hs")
    foldM withFile acc files

withFile :: Map.Map Char Int -> FilePath -> IO (Map.Map Char Int)
withFile acc fp = do
    contents <- BS.readFile fp
    let str = fromUTF8BS contents
    e <- parseModuleEx defaultParseFlags fp (Just str)
    res <- case e of
        Left _  -> return Nothing
        Right x -> return (Just (ghcModule x))

    let m = SYB.everything (Map.unionWith (+)) (SYB.mkQ Map.empty withExpr) res
    return $! Map.unionWith (+) acc m

withExpr :: HsExpr GhcPs -> Map.Map Char Int
withExpr (HsVar _ (L _ (Unqual o)))
    | [c] <- occNameString o
    , 'a' <= c
    , c <= 'z'
    = Map.singleton c 1

withExpr _ = Map.empty