ExtractTopLevel.hs

-- | This code is from an FP Complete internal tool, and isn't
-- necessarily fit for packaging.  It's released under the MIT
-- license.
module ExtractTopLevel (main) where
 
import           Data.Char (isSpace)
import           Data.Generics.Schemes (listify)
import           Data.List (find)
import qualified Language.Haskell.Exts as Simple
import           Language.Haskell.Exts.Annotated
import           Language.Haskell.Exts.Annotated.Simplify (sExp, sName)
 
main :: IO ()
main = do
    let parseMode = defaultParseMode { fixities = Nothing }
        filePath = "SomeFile.hs"
        varName = Simple.Ident "someIdentifier"
    result <- parseFileWithComments parseMode filePath
    case result of
        ParseOk (Module _ _ _ _ decls, comments) ->
            case find (funcNameMatches varName) decls of
                Just decl -> putStrLn $ exactPrintPortion comments decl
                Nothing -> fail "Didn't find function"
        _ -> fail (show result)
 
type SSI = SrcSpanInfo
 
exactPrintPortion :: (ExactP a, Annotated a) => [Comment] -> a SSI -> String
exactPrintPortion comments x =
    reverse $ dropWhile (`elem` "\r\n") $ reverse $ -- dropWhileEnd
    dropWhile (`elem` "\r\n") $
    exactPrint x containedComments
  where
    ss = srcInfoSpan (ann x)
    -- This is needed because some declaration source spans extend
    -- until the beginning of the next declaration.
    lastLine =
        length $
        dropWhile (all isSpace) $
        reverse $
        lines $
        exactPrint x []
    containedComments =
        [ comment
        | comment@(Comment _ cs _) <- comments
        , ssStart cs >= ssStart ss
        , ssEnd cs <= (lastLine + 1, 0)
        ]
 
-- Use these to find top level decls by name
 
funcNameMatches :: Simple.Name -> Decl SSI -> Bool
funcNameMatches n (FunBind _ (Match _ n' _ _ _ : _)) = n == sName n'
funcNameMatches n (PatBind _ pat _ _) = not $ null $ listify (pvarMatches n) pat
funcNameMatches n (ForImp _ _ _ _ n' _) = n == sName n'
funcNameMatches n (ForExp _ _ _ n' _) = n == sName n'
funcNameMatches _ _ = False
 
sigMatches :: Simple.Name -> Decl SSI -> Bool
sigMatches n (TypeSig _ ns _) = any ((== n) . sName) ns
sigMatches _ _ = False
 
typeMatches :: Simple.Name -> Decl SSI -> Bool
typeMatches n = maybe False (declHeadMatches n) . getDeclHead
 
-- Utilities
 
pvarMatches :: Simple.Name -> Pat SSI -> Bool
pvarMatches n (PVar _ n') = n == sName n'
pvarMatches _ _ = False
 
declHeadMatches :: Simple.Name -> DeclHead SSI -> Bool
declHeadMatches n (DHead _ n') = n == sName n'
declHeadMatches n (DHInfix _ _ n') = n == sName n'
declHeadMatches n (DHParen _ dh) = declHeadMatches n dh
declHeadMatches n (DHApp _ dh _) = declHeadMatches n dh
 
getDeclHead :: Decl SSI -> Maybe (DeclHead SSI)
getDeclHead (TypeDecl _ dh _) = Just dh
getDeclHead (TypeFamDecl _ dh _) = Just dh
getDeclHead (ClosedTypeFamDecl _ dh _ _) = Just dh
getDeclHead (DataDecl _ _ _ dh _ _) = Just dh
getDeclHead (GDataDecl _ _ _ dh _ _ _) = Just dh
getDeclHead (DataFamDecl _ _ dh _) = Just dh
getDeclHead (ClassDecl _ _ dh _ _) = Just dh
getDeclHead _ = Nothing
 
ssStart :: SrcSpan -> (Int, Int)
ssStart ss = (srcSpanStartLine ss, srcSpanStartColumn ss)
 
ssEnd :: SrcSpan -> (Int, Int)
ssEnd ss = (srcSpanEndLine ss, srcSpanEndColumn ss)