1 files changed, 13 insertions, 12 deletions

diff --git a/edit-lens/src/Data/String/DFST.hs b/edit-lens/src/Data/String/DFST.hs
index 8a22dd3..75ca1cd 100644
--- a/edit-lens/src/Data/String/DFST.hs
+++ b/edit-lens/src/Data/String/DFST.hs
@@ -8,7 +8,7 @@ Description: Deterministic finite state transducers
 -}
 module Data.String.DFST
   ( DFST(..)
-  , runDFST
+  , runDFST, runDFST'
   ) where
 
 import Data.Map.Strict (Map, (!?))
@@ -17,6 +17,8 @@ import qualified Data.Map.Strict as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
 
+import Control.Monad
+
 data DFST state = DFST
   { stInitial :: state
   , stTransition :: Map (state, Char) (state, Maybe Char)
@@ -25,14 +27,13 @@ data DFST state = DFST
   }
   
 runDFST :: forall state. Ord state => DFST state -> String -> Maybe String
-runDFST DFST{..} str = ($ []) <$> go stInitial str id
-  where
-    go :: state -> String -> (String -> String) -> Maybe (String -> String)
-    go st [] acc
-      | st `Set.member` stAccept = Just acc
-      | otherwise = Nothing
-    go st (c:cs) acc
-      | Just (st', mc') <- stTransition !? (st, c)
-      = go st' cs $ acc . maybe id (:) mc'
-      | otherwise
-      = go st cs acc
+runDFST dfst@DFST{..} str = let (finalState, str') = runDFST' dfst stInitial str id
+                            in str' "" <$ guard (finalState `Set.member` stAccept)
+
+runDFST' :: forall state. Ord state => DFST state -> state -> String -> (String -> String) -> (state, (String -> String))
+runDFST' _ st [] acc = (st, acc)
+runDFST' dfst@DFST{..} st (c:cs) acc
+  | Just (st', mc') <- stTransition !? (st, c)
+  = runDFST' dfst st' cs $ acc . maybe id (:) mc'
+  | otherwise
+  = runDFST' dfst st cs acc