1 files changed, 13 insertions, 7 deletions

diff --git a/edit-lens/src/Control/DFST.lhs b/edit-lens/src/Control/DFST.lhs
index aec7bbb..476a8c4 100644
--- a/edit-lens/src/Control/DFST.lhs
+++ b/edit-lens/src/Control/DFST.lhs
@@ -39,25 +39,31 @@ data DFST state input output = DFST
   }
 
 
-toFST :: forall state input output. (Ord state, Ord input, Ord output) => DFST state input output -> FST (state, Natural) input output
+toFST :: forall state input output. (Ord state, Ord input, Ord output) => DFST state input output -> FST (state, Maybe (input, Natural)) input output
 -- ^ Split apart non-singleton outputs into a series of epsilon-transitions
 --
 -- This function is currently invalid since multiple out-edges in the `DFST` visit the same intermediate states (we should label intermediate states not only with an ordinal but also with the input Symbol from the `DFST`)
 toFST DFST{..} = flip execState initialFST $ forM_ (Map.toList stTransition) handleTransition
   where
     initialFST = FST
-      { stInitial = (stInitial, 0)
+      { stInitial = (stInitial, Nothing)
       , stTransition = Map.empty
-      , stAccept = Set.map (,0) stAccept
+      , stAccept = Set.map (, Nothing) stAccept
       }
     addTransition :: forall state input output. (Ord state, Ord input, Ord output) => (state, Maybe input) -> (state, Maybe output) -> State (FST state input output) ()
     addTransition k v = modify $ \f@FST{ stTransition } -> f { FST.stTransition = Map.insertWith Set.union k (Set.singleton v) stTransition }
-    handleTransition :: ((state, input), (state, Seq output)) -> State (FST (state, Natural) input output) ()
-    handleTransition ((st, inS), (st', outs)) = handleTransition' (st, 0) (Just inS) outs (st', 0)
-    handleTransition' :: (state, Natural) -> Maybe input -> Seq output -> (state, Natural) -> State (FST (state, Natural) input output) ()
+    handleTransition :: ((state, input), (state, Seq output)) -> State (FST (state, Maybe (input, Natural)) input output) ()
+    handleTransition ((st, inS), (st', outs)) = handleTransition' (st, Nothing) (Just inS) outs (st', Nothing)
+    handleTransition' :: (state, Maybe (input, Natural)) -> Maybe input -> Seq output -> (state, Maybe (input, Natural)) -> State (FST (state, Maybe (input, Natural)) input output) ()
     handleTransition' from inS Empty to = addTransition (from, inS) (to, Nothing)
     handleTransition' from inS (outS :<| Empty) to = addTransition (from, inS) (to, Just outS)
-    handleTransition' from@(st, i) inS (outS :<| oo) to = addTransition (from, inS) ((st, succ i), Just outS) >> handleTransition' (st, succ i) Nothing oo to
+    handleTransition' from@(st, chain) inS (outS :<| oo) to = do
+      let next
+            | Just (inS', i) <- chain = (st, Just (inS', succ i))
+            | Just inS'      <- inS   = (st, Just (inS', 0     ))
+            | otherwise               = error "TODO: Can this happen?"
+      addTransition (from, inS) (next, Just outS)
+      handleTransition' next Nothing oo to
   
 runDFST :: forall state input output. (Ord state, Ord input) => DFST state input output -> Seq input -> Maybe (Seq output)
 runDFST dfst@DFST{..} str = let (finalState, str') = runDFST' dfst stInitial str Seq.empty