pretty print the AST generated by Language.Haskell.Exts

hse-pprint.hs

{- Simple pretty printer for the AST created by Language.Haskell.Exts.
 -
 - Usage: program source.hs
 -
 -}

import qualified Language.Haskell.Exts as E
import Language.Haskell.Exts.Parser (ParseResult(..))
import Data.List
import System.Environment
import System.IO
import Control.Monad
import Text.ParserCombinators.Parsec
import Text.ParserCombinators.Parsec.Expr
import Text.ParserCombinators.Parsec.Language
import qualified Text.ParserCombinators.Parsec.Token as Token

spaces1 :: Parser ()
spaces1 = skipMany1 space

data Sexp
  = Atom String
  | Str  String
  | Chr  Char
  | Seq Char Char [Sexp]
  deriving (Show)

languageDef = haskellDef

lexer = Token.makeTokenParser languageDef

identifier = Token.identifier lexer -- parses an identifier
reserved   = Token.reserved   lexer -- parses a reserved name
reservedOp = Token.reservedOp lexer -- parses an operator
parens     = Token.parens     lexer -- parses surrounding parenthesis:
                                    --   parens p
                                    -- takes care of the parenthesis and
                                    -- uses p to parse what's inside them
integer    = Token.integer    lexer -- parses an integer
semi       = Token.semi       lexer -- parses a semicolon
whiteSpace = Token.whiteSpace lexer -- parses whitespace

parseSexp = do
  whiteSpace 
  try parseAtom <|> try parseString <|> try parseChar <|> parseSeq

parseString = fmap Str $ Token.stringLiteral lexer
parseChar = fmap Chr $ Token.charLiteral lexer

{-
parseString = do
  char '"'
  s <- many1 $ oneOf "abcdefgh123zs"
  char '"'
  return $ Str s
-}

parseAtom = fmap Atom $ many1 $ noneOf " \t\r\n()[]{}'\""

{-
parseAtom = (fmap Atom $ Token.identifier lexer)
            <|> (fmap (Atom . show) $ Token.integer lexer)
-}

parenMate :: Char -> Char
parenMate c = case c of
                '(' -> ')'
                '[' -> ']'
                '{' -> '}'
                _   -> undefined

parseSeq = do
  c <- oneOf "([{"
  let d = parenMate c
  xs <- sepEndBy parseSexp spaces
  char d
  return $ Seq c d xs

sexpLength :: Sexp -> Int
sexpLength (Atom s) = length s
sexpLength (Str s) = length $ show s
sexpLength (Chr c)   = length $ show c
sexpLength (Seq c d items) = 1 + sum (map sexpLength items) + length items

emit depth str = putStrLn $ (replicate (2*depth) ' ') ++ str

sexpToString (Atom s) = s
sexpToString (Str s)  = show s
sexpToString (Chr c)  = show c
sexpToString (Seq c d items)
  = [c] ++ intercalate " " (map sexpToString items) ++ [d]

singleLine :: Int -> [Sexp] -> Maybe String
singleLine limit items = go limit items []
  where
    go limit [] acc = Just $ if null acc then [] else tail acc
    go limit ((Seq _ _ _):xs) acc = Nothing
    go limit (x:xs) acc
      | limit >= length s + 1 = go limit' xs  (acc ++ " " ++ s)
      | otherwise             = Nothing
      where s = sexpToString x
            limit' = limit - (length s - 1)

test1 = singleLine 80 []
test2 = singleLine 80 [Atom "anc"]

fmtSexp :: Int -> Sexp -> IO ()
fmtSexp depth seq@(Seq c d []) = do
  emit depth $ [c] ++ " " ++ [d]
fmtSexp depth seq@(Seq c d items@(x:xs)) = do
  case singleLine 80 items of
    Just str -> emit depth $ [c] ++ str ++ [d]
    Nothing  -> do case singleLine 80 [x] of
                     Just str -> do emit depth $ [c] ++ " " ++ str
                                    forM_ xs $ fmtSexp (depth+1)
                     Nothing  -> do emit depth [c]
                                    forM_ items $ fmtSexp (depth+1)
                   emit depth [d]
fmtSexp depth (Atom s) = emit depth s
fmtSexp depth (Str s) = emit depth (show s)
fmtSexp depth (Chr c) = emit depth (show c)

parseHaskell path = do
  r <- E.parseFile path
  case r of
    ParseOk v -> return v
    _         -> error $ "unable to parse Haskell source in " ++ path

parseAst tree = 
  case parse parseSexp "exp"  tree of
    Left err -> error $ "problems parsing AST: " ++ show err
    Right val -> val

main = do
  (path:_) <- getArgs
  ast <- parseHaskell path
  let tree = "(" ++ show ast ++ ")"
      sexp = parseAst tree
  fmtSexp 0 sexp