path: root/src/Main.hs

module Main where

import Control.Monad
import Data.List
import Data.Maybe
import Haskell.AST
import Haskell.Env
import Haskell.Rewrite
import Haskell.Parser
import Pretty
import System.Environment
import System.Exit
import System.IO


tryEither :: Either String a -> IO a
tryEither (Left err) = die err
tryEither (Right x) = return x

tryEither' :: Show e => Either e a -> IO a
tryEither' (Left err) = die (show err)
tryEither' (Right x) = return x

fromRight :: Either a b -> b
fromRight (Right x) = x
fromRight (Left _) = error "fromRight on Left"

isRight :: Either a b -> Bool
isRight (Right _) = True
isRight (Left _) = False

while :: Monad m => a -> (a -> m (Maybe a)) -> m ()
while val f = f val >>= \case
    Nothing   -> return ()
    Just val' -> while val' f

whenM :: Monad m => m Bool -> m a -> m ()
whenM b a = b >>= \b' -> if b' then void a else return ()

ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM b t e = b >>= \b' -> if b' then t else e

splitOn :: Eq a => a -> [a] -> [[a]]
splitOn _   [] = [[]]
splitOn spl (x:xs)
    | x == spl  = [] : splitOn spl xs
    | otherwise = let (r : rs) = splitOn spl xs in (x:r) : rs


usageString :: String
usageString =
    "\x1B[1mMeta-commands\x1B[0m\n\
    \  env          -- Show the current environment\n\
    \  exit         -- Exit program immediately\n\
    \  undo         -- Undo last operation\n\
    \  focus \x1B[4mx\x1B[0m      -- Watch and operate on given variable\n\
    \  focus        -- Unfocus currently focused variable\n\
    \  log          -- Print list of commands that reproduces current environment\n\
    \  help         -- Show this help\n\
    \  auto         -- Try to simplify the currently focused expression as far as possible\n\
    \  rewall       -- Rewrite as much as possible, but each name at most once\n\
    \\n\
    \\x1B[1mExpression rewrite commands\x1B[0m\n\
    \These commands operate on the currently-focused expression.\n\
    \  rew \x1B[4mx\x1B[0m        -- Replaces occurrences of \x1B[4mx\x1B[0m with its value\n\
    \  beta         -- Applies some beta-reductions ((\\x -> M) y => M[y/x])\n\
    \  eta          -- Applies some eta-reductions ((\\x -> f x) => f)\n\
    \  case         -- Evaluates case statements\n\
    \  case!        -- Evaluates case statements, choosing first hit if multiple branches unify\n\
    \  etacase      -- If all case branches are the same as their patterns, remove case statement\n\
    \  casecase     -- case (case x of {A->a;B->b}) of {C->c;D->d} \n\
    \                    => case x of {A->case a of {C->c;D->d}; B->case b of {C->c;D->d}}\n\
    \  repeat \x1B[4mcmd\x1B[0m   -- Repeat \x1B[4mcmd\x1B[0m until no change occurs\n\
    \\n\
    \To evaluate a rewrite command on an unfocused variable, use:\n\
    \  action \x1B[4mx\x1B[0m \x1B[4mcmd\x1B[0m -- Evaluate \x1B[4mcmd\x1B[0m on variable \x1B[4mx\x1B[0m"


data AppState = AppState { asCtx :: Context, asFocus :: Maybe Name }
  deriving (Show)

data Context = Context { cStack :: [(Action, Env)], cBaseEnv :: Env }
  deriving (Show)

data Action = Action Cmd Name
  deriving (Show)

data Cmd = CRewrite Name
         | CBeta
         | CEta
         | CCase
         | CCaseForce
         | CEtaCase
         | CCaseCase
         | CRepeat Cmd
  deriving (Show)

data UserCmd = UCCmd Cmd
             | CAction Action
             | CAuto
             | CRewriteAll
             | CExit
             | CUndo
             | CShowEnv
             | CUnFocus
             | CFocus Name
             | CLog
             | CHelp
  deriving (Show)

instance Pretty Action where
    pretty (Action cmd target) = Node ("action " ++ target) [pretty cmd]

instance Pretty Cmd where
    pretty (CRewrite name) = Leaf ("rew " ++ name)
    pretty CBeta = Leaf "beta"
    pretty CEta = Leaf "eta"
    pretty CCase = Leaf "case"
    pretty CCaseForce = Leaf "case!"
    pretty CEtaCase = Leaf "etacase"
    pretty CCaseCase = Leaf "casecase"
    pretty (CRepeat cmd) = Node "repeat" [pretty cmd]

topEnv :: Context -> Env
topEnv (Context [] env) = env
topEnv (Context ((_, env):_) _) = env

pushCtx :: (Action, Env) -> Context -> Context
pushCtx pair ctx = ctx { cStack = pair : cStack ctx }

exprTransformer :: Cmd -> Maybe (Expr -> Expr)
exprTransformer CBeta = Just betared
exprTransformer CEta = Just etared
exprTransformer CCase = Just (casered False)
exprTransformer CCaseForce = Just (casered True)
exprTransformer CEtaCase = Just etacase
exprTransformer CCaseCase = Just casecase
exprTransformer _ = Nothing

ctxTransform :: Context -> Action -> (Expr -> Expr) -> Either String Context
ctxTransform ctx act@(Action _ target) = generalCtxTransform ctx act (\env -> get env target)

generalCtxTransform :: Context -> Action -> (Env -> Either String a) -> (a -> Expr) -> Either String Context
generalCtxTransform ctx act@(Action _ target) getter f = do
    expr <- f <$> getter (topEnv ctx)
    env' <- reAdd (topEnv ctx) (Def target expr)
    return $ pushCtx (act, env') ctx

apply :: Context -> Action -> Either String Context
apply ctx act@(Action cmd target) = case cmd of
    CRewrite name ->
        generalCtxTransform
            ctx act
            (\env -> liftM2 (,) (get env name) (get env target))
            (\(repl, expr) -> normalise $ rewrite name repl expr)
    CRepeat subcmd ->
        case exprTransformer subcmd of
            Nothing -> Left $ "Cannot repeat '" ++ pprint 80 subcmd ++ "'"
            Just trans -> ctxTransform ctx act (fixpoint (normalise . trans))
    _ ->
        ctxTransform ctx act (normalise . fromJust (exprTransformer cmd))

findChangingTransformer :: Expr -> Maybe Cmd
findChangingTransformer expr =
    let opts = [CBeta, CEta, CCase, CEtaCase, CCaseCase]
        results = [(cmd, normalise $ fromJust (exprTransformer cmd) expr) | cmd <- opts]
    in fst <$> find (\(_, res) -> res /= expr) results

applyUserCmd :: AppState -> UserCmd -> IO (Either String (Maybe AppState))
applyUserCmd appstate = \case
    UCCmd cmd ->
        case asFocus appstate of
            Nothing -> return (Left "No focus set")
            Just focus ->
                case apply (asCtx appstate) (Action cmd focus) of
                    Left err -> return (Left err)
                    Right ctx -> return (Right (Just appstate { asCtx = ctx }))
    CAction act ->
        case apply (asCtx appstate) act of
            Left err -> return (Left err)
            Right ctx -> return (Right (Just appstate { asCtx = ctx }))
    CAuto ->
        case asFocus appstate of
            Nothing -> return (Left "No focus set")
            Just focus ->
                let nextStep ctx = case findChangingTransformer (fromRight (get (topEnv ctx) focus)) of
                        Nothing  -> Right ctx
                        Just cmd -> case apply ctx (Action cmd focus) of
                            Left err   -> Left err
                            Right ctx' -> nextStep ctx'
                in return $ fmap (\ctx -> Just appstate { asCtx = ctx }) $ nextStep (asCtx appstate)
    CRewriteAll ->
        case asFocus appstate of
            Nothing -> return (Left "No focus set")
            Just focus ->
                let nextStep ctx done =
                        let orig = fromRight (get (topEnv ctx) focus)
                        in case filter (isRight . get (topEnv ctx)) $ allRefs orig \\ done of
                            [] -> Right ctx
                            (name:_) -> case apply ctx (Action (CRewrite name) focus) of
                                Left err   -> Left ("Error rewriting '" ++ name ++ "': " ++ err)
                                Right ctx' ->
                                    let newexpr = fromRight (get (topEnv ctx') focus)
                                    in nextStep (if newexpr /= orig then ctx' else ctx) (name : done)
                in return $ fmap (\ctx -> Just appstate { asCtx = ctx }) $ nextStep (asCtx appstate) []
    CExit ->
        return (Right Nothing)
    CUndo ->
        let ctx = asCtx appstate
        in case cStack ctx of
                 [] -> return (Left "Empty environment stack")
                 (_:stk) -> return (Right (Just appstate { asCtx = ctx { cStack = stk } }))
    CShowEnv ->
        putStrLn (pprint 80 (topEnv (asCtx appstate))) >> return (Right (Just appstate))
    CUnFocus ->
        return (Right (Just appstate { asFocus = Nothing }))
    CFocus name ->
        if envContains (topEnv (asCtx appstate)) name
            then return (Right (Just appstate { asFocus = Just name }))
            else return (Left "Name doesn't exist")
    CLog -> do
        mapM_ (putStrLn . pprint 80) $ map fst (reverse $ cStack (asCtx appstate))
        return (Right (Just appstate))
    CHelp ->
        putStrLn usageString >> return (Right (Just appstate))

applyUserCmds :: AppState -> [UserCmd] -> IO (Either String (Maybe AppState))
applyUserCmds appstate [] = return (Right (Just appstate))
applyUserCmds appstate (ucmd : ucmds) =
    applyUserCmd appstate ucmd >>= \case
        Left err -> return (Left err)
        Right Nothing -> return (Right Nothing)
        Right (Just as) -> applyUserCmds as ucmds

interface :: AppState -> IO (Maybe AppState)
interface appstate = do
    putStrLn ""

    let appstate' = case asFocus appstate of
                        Nothing -> appstate
                        Just focus ->
                            if envContains (topEnv (asCtx appstate)) focus
                                then appstate
                                else appstate { asFocus = Nothing }

    case asFocus appstate' of
        Nothing -> return ()
        Just n  ->
            case get (topEnv (asCtx appstate')) n of
                Left _ -> return ()
                Right expr -> putStrLn $ pprint 80 (Def n expr)

    putStr "> " >> hFlush stdout

    ifM isEOF (return Nothing)
              (getLine >>= handleInputLine appstate')

handleInputLine :: AppState -> String -> IO (Maybe AppState)
handleInputLine appstate line = do
    let phrases = map words $ splitOn ';' line

    ucmds <- case mapM parseCmd $ filter (not . null) phrases of
                Left err -> putStrLn err >> return []
                Right ucmds -> return ucmds

    applyUserCmds appstate ucmds >>= \case
        Left err -> putStrLn err >> return (Just appstate)
        Right Nothing -> return Nothing
        Right (Just as) -> return (Just as)

parseCmd :: [String] -> Either String UserCmd
parseCmd ("action" : target : rest) = parseCmd rest >>= \case
    UCCmd cmd -> Right (CAction (Action cmd target))
    _ -> Left "Invalid command within 'action'"
parseCmd ["exit"] = Right CExit
parseCmd ["undo"] = Right CUndo
parseCmd ["env"] = Right CShowEnv
parseCmd ["focus"] = Right CUnFocus
parseCmd ["focus", name] = Right (CFocus name)
parseCmd ["log"] = Right CLog
parseCmd ["help"] = Right CHelp
parseCmd ["auto"] = Right CAuto
parseCmd ["rewall"] = Right CRewriteAll
parseCmd ["rew", func] = Right (UCCmd (CRewrite func))
parseCmd ["beta"] = Right (UCCmd CBeta)
parseCmd ["eta"] = Right (UCCmd CEta)
parseCmd ["case"] = Right (UCCmd CCase)
parseCmd ["case!"] = Right (UCCmd CCaseForce)
parseCmd ["etacase"] = Right (UCCmd CEtaCase)
parseCmd ["casecase"] = Right (UCCmd CCaseCase)
parseCmd ("repeat" : rest) = parseCmd rest >>= \case
    UCCmd (CRepeat _) -> Left "Cannot repeat 'repeat'"
    UCCmd cmd -> Right (UCCmd (CRepeat cmd))
    _ -> Left "Invalid command within 'repeat'"
parseCmd cmd = Left $ "Unrecognised command: " ++ show cmd

main :: IO ()
main = do
    fname <- getArgs >>= \case
        [fname] -> return fname
        _       -> die "Usage: verify-hs <file.hs>"

    source <- readFile fname
    ast <- tryEither' (parseAST fname source)
    print ast
    env0 <- tryEither (envFromAST ast)
    print env0

    while (AppState (Context [] env0) Nothing) interface