Repeat commands; add functor proof

1 files changed, 28 insertions, 8 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 3ed39f9..bdbd447 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,6 +1,7 @@
 module Main where
 
 import Control.Monad
+import Data.Maybe
 import Haskell.AST
 import Haskell.Env
 import Haskell.Rewrite
@@ -53,6 +54,7 @@ data Cmd = CForget
          | CCaseForce
          | CEtaCase
          | CCaseCase
+         | CRepeat Cmd
   deriving (Show)
 
 data UserCmd = UCCmd Cmd
@@ -77,6 +79,7 @@ instance Pretty Cmd where
     pretty CCaseForce = "case!"
     pretty CEtaCase = "etacase"
     pretty CCaseCase = "casecase"
+    pretty (CRepeat cmd) = "repeat " ++ pretty cmd
 
 topEnv :: Context -> Env
 topEnv (Context [] env) = env
@@ -85,26 +88,39 @@ 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
+
+repeatTrans :: (Expr -> Expr) -> Expr -> Expr
+repeatTrans f expr =
+    let values = iterate f expr
+        pairs = zip values (tail values)
+    in fst . head $ dropWhile (uncurry (/=)) pairs
+
 apply :: Context -> Action -> Either String Context
 apply ctx act@(Action cmd target) = case cmd of
     CForget -> do
         env' <- forget (topEnv ctx) target
         return $ pushCtx (act, env') ctx
     CRewrite name -> genTransform (\env -> liftM2 (,) (get env name) (get env target))
-                                  (\(repl, expr) -> rewrite name repl expr)
-    CBeta -> transform betared
-    CEta -> transform etared
-    CCase -> transform (casered False)
-    CCaseForce -> transform (casered True)
-    CEtaCase -> transform etacase
-    CCaseCase -> transform casecase
+                                  (\(repl, expr) -> normalise $ rewrite name repl expr)
+    CRepeat subcmd -> case exprTransformer subcmd of
+                          Nothing -> Left $ "Cannot repeat '" ++ pretty subcmd ++ "'"
+                          Just trans -> transform (repeatTrans (normalise . trans))
+    _ -> transform (normalise . fromJust (exprTransformer cmd))
   where
     transform :: (Expr -> Expr) -> Either String Context
     transform = genTransform (\env -> get env target)
 
     genTransform :: (Env -> Either String a) -> (a -> Expr) -> Either String Context
     genTransform getter f = do
-        expr <- normalise . f <$> getter (topEnv ctx)
+        expr <- f <$> getter (topEnv ctx)
         env' <- reAdd (topEnv ctx) (Def target expr)
         return $ pushCtx (act, env') ctx
 
@@ -197,6 +213,10 @@ 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 ()