1 files changed, 51 insertions, 2 deletions

diff --git a/edit-lens/src/Control/FST.lhs b/edit-lens/src/Control/FST.lhs
index 4f1f364..9298e11 100644
--- a/edit-lens/src/Control/FST.lhs
+++ b/edit-lens/src/Control/FST.lhs
@@ -24,11 +24,11 @@ import qualified Data.Set as Set
 import Data.Sequence (Seq)
 import qualified Data.Sequence as Seq
 
-import Data.Maybe (mapMaybe, fromMaybe)
+import Data.Maybe (mapMaybe, fromMaybe, isJust, fromJust)
 
 import Numeric.Natural
 
-import Control.Lens.TH
+import Control.Lens
 
 import Control.Monad.State.Strict
 
@@ -59,6 +59,55 @@ instance (Show state, Show input, Show output) => Pretty (FST state input output
       list :: [PP.Doc] -> PP.Doc
       list = PP.encloseSep (PP.lbracket PP.<> PP.space) (PP.space PP.<> PP.rbracket) (PP.comma PP.<> PP.space)
 
+runFST :: forall input output state. (Ord input, Ord output, Ord state) => FST state input output -> Seq input -> [Seq output]
+runFST = fmap (map $ catMaybes . fmap (view _2) . view _2) . runFST'
+  where
+    catMaybes = fmap fromJust . Seq.filter isJust
+
+runFST' :: forall input output state. (Ord input, Ord output, Ord state)
+        => FST state input output
+        -> Seq input
+        -> [(state, Seq (state, Maybe output))] -- ^ Tuples of initial state and chosen transitions; not neccessarily finite
+-- ^ Compute all possible runs on the given input
+runFST' fst Seq.Empty = guardAccept $ (\(_, st, _) -> (st, Seq.Empty)) <$> step fst Nothing Nothing
+runFST' fst cs = guardAccept $ do
+  initial <- view _2 <$> step fst Nothing Nothing
+  go (initial, Seq.Empty) cs
+  where
+    guardAccept res = do
+      (initial, path) <- res
+      let
+        finalState
+          | (_ :> (st, _)) <- path = st
+          | otherwise              = initial
+      guard $ finalState `Set.member` stAccept
+      return res
+
+    go :: (state, Seq (state, Maybe output)) -> Seq input-> [(state, Seq (state, Maybe output))]
+    go (initial, path) cs = do
+      let
+        current
+          | (_ :> (st, _)) <- path = st
+          | otherwise              = initial
+      (head, next, out) <- step fst (Just current) (Seq.lookup 0 cs)
+      let
+        nPath = path :> (next, out)
+        ncs = maybe id (:<) head cs
+      go (initial, nPath) ncs
+      
+
+step :: forall input output state. (Ord input, Ord output, Ord state)
+     => FST state input output
+     -> Maybe state -- ^ Current state
+     -> Maybe input -- ^ Head of remaining input
+     -> [(Maybe input, state, Maybe output)] -- ^ Tuples of unconsumed input, next state, and produced output
+step FST{..} Nothing inS = (\s -> (inS, s, Nothing)) <$> Set.toList stInitial
+step FST{..} (Just c) inS = let
+  consuming = fromMaybe Set.empty $ Map.lookup (c, inS) stTransition
+  unconsuming = fromMaybe Set.empty $ Map.lookup (c, Nothing) stTransition
+  in Set.toList $ Set.map (\(n, mOut) -> (Nothing, n, mOut)) consuming `Set.union` Set.map (\(n, mOut) -> (inS, n, mOut)) unconsuming
+  
+
 wordFST :: forall input output. Seq output -> FST Natural input output
 -- ^ @wordFST str@ is the linear FST generating @str@ as output when given no input
 wordFST outs = FST