Standalone version of lambdabot's undo function.

undo.hs

-- Copyright (c) 2006 Spencer Janssen
-- GPL version 2 or later (see http://www.gnu.org/copyleft/gpl.html)

{-# LANGUAGE Rank2Types #-}

module Main where


import Control.Monad
import Data.Generics
import qualified Data.Set as Set
import Language.Haskell.Exts.Syntax hiding (Module)
import Data.Generics
import System.Environment (getArgs)
import Language.Haskell.Exts
	( parseExp
	, parseDecl
	, ParseResult(ParseOk, ParseFailed)
	, Pretty
	, prettyPrint
	, prettyPrintWithMode
	, defaultMode
	, layout
	, PPLayout(PPInLine)
	)

main = interactArgsOrInput (transform undo)

interactArgs :: (String -> String) -> IO ()
interactArgs f = getArgs >>= (putStrLn . f . unwords)

interactArgsOrInput :: (String -> String) -> IO ()
interactArgsOrInput f = getArgs >>= (\args -> pickInteract (length args) f)
	where
		pickInteract :: Int -> (String -> String) -> IO ()
		pickInteract argCount = if (argCount == 0) then interact else interactArgs

-- |Parse a string as an 'Exp' or a 'Decl', apply the given generic transformation to it,
-- and re-render it back to text.
withParsed :: (forall a. (Data a, Eq a) => a -> a) -> String -> String
withParsed _ "" = "Error: expected a Haskell expression or declaration"
withParsed f s  = case (parseExp s, parseDecl s) of
    (ParseOk a, _)          -> prettyPrintInLine $ f a
    (_, ParseOk a)          -> prettyPrintInLine $ f a
    (ParseFailed l e,  _)   -> prettyPrint l ++ ':' : e

-- |Render haskell code in a compact format
prettyPrintInLine :: Pretty a => a -> String
prettyPrintInLine = prettyPrintWithMode (defaultMode { layout = PPInLine })


findVar :: Data a => a -> String
findVar e = head $ do
    i <- [0 ..]
    x <- ['a' .. 'z']
    let xi = x : replicate i '\''
    guard $ not $ Set.member xi s
    return xi
 where s = Set.fromList $ listify (const True :: String -> Bool) e

transform :: (String -> Exp -> Exp) -> String -> String
transform f = withParsed $ \e -> everywhere (mkT . f . findVar $ e) e

undo :: String -> Exp -> Exp
undo v (Do stms) = f stms
 where
    f [Qualifier e]          = e
    f (Qualifier e     : xs) = infixed e ">>" $ f xs
    f (LetStmt   ds    : xs) = Let ds $ f xs
    f (Generator s p e : xs)
        | irrefutable p = infixed e ">>=" $ Lambda s [p] $ f xs
        | otherwise     = infixed e ">>=" $
                            Lambda s [pvar v] $
                                Case (var v)
                                    [ alt p (f xs)
                                    , alt PWildCard $
                                        App
                                            (var "fail")
                                            (Lit $ String "")
                                    ]
        where alt pat x = Alt s pat (UnGuardedRhs x) (BDecls [])
    f _ = error "Undo plugin error: can't undo!"
undo v (ListComp e stms) = f stms
 where
    f []                                = List [e]
    f (QualStmt (Qualifier g    ) : xs) = If g (f xs) nil
    f (QualStmt (LetStmt   ds   ) : xs) = Let ds $ f xs
    f (QualStmt (Generator s p l) : xs)
        | irrefutable p = concatMap' $ Lambda s [p] $ f xs
        | otherwise     = concatMap' $
                            Lambda s [pvar v] $
                                Case (var v)
                                    [ alt p (f xs)
                                    , alt PWildCard nil
                                    ]
        where alt pat x = Alt s pat (UnGuardedRhs x) (BDecls [])
              concatMap' fun = App (App (var "concatMap") (Paren fun)) l
    f _ = error "Undo plugin error: can't undo!"
undo _ x           = x

irrefutable :: Pat -> Bool
irrefutable (PVar _)     = True
irrefutable (PIrrPat _)  = True
irrefutable PWildCard    = True
irrefutable (PAsPat _ p) = irrefutable p
irrefutable (PParen p)   = irrefutable p
irrefutable (PTuple _box ps) = all irrefutable ps
irrefutable _              = False

infixed :: Exp -> String -> Exp -> Exp
infixed l o r = InfixApp l (QVarOp $ UnQual $ Symbol o) r

nil :: Exp
nil = Var list_tycon_name

var :: String -> Exp
var = Var . UnQual . Ident

pvar :: String -> Pat
pvar = PVar . Ident

do' :: String -> Exp -> Exp
do' _ (Let ds (Do s)) = Do (LetStmt ds : s)
do' v e@(InfixApp l (QVarOp (UnQual (Symbol op))) r) =
     case op of
         ">>=" ->
             case r of
                 (Lambda loc [p] (Do stms)) -> Do (Generator loc p l : stms)
                 (Lambda loc [PVar v1] (Case (Var (UnQual v2))
                                            [ Alt _ p (UnGuardedRhs s) (BDecls [])
                                            , Alt _ PWildCard (UnGuardedRhs (App (Var (UnQual (Ident "fail"))) _)) (BDecls [])
                                            ]))
                           | v1 == v2           -> case s of
                                                       Do stms -> Do (Generator loc p l : stms)
                                                       _         -> Do [Generator loc p l, Qualifier s]
                 (Lambda loc [p] s)           -> Do [Generator loc p l, Qualifier s]
                 _ -> Do [ Generator undefined (pvar v) l
                           , Qualifier . app r $ var v]
         ">>" ->
             case r of
                 (Do stms) -> Do (Qualifier l : stms)
                 _           -> Do [Qualifier l, Qualifier r]
         _    -> e
do' _ x = x

-- | 'app' is a smart constructor that inserts parens when the first argument
-- is an infix application.
app :: Exp -> Exp -> Exp
app e@(InfixApp {}) f = App (Paren e) f
app e                 f = App e f