3 files changed, 114 insertions, 3 deletions

diff --git a/edit-lens/package.yaml b/edit-lens/package.yaml
index 2374898..ec10b83 100644
--- a/edit-lens/package.yaml
+++ b/edit-lens/package.yaml
@@ -33,6 +33,7 @@ library:
     - composition-tree
     - Diff
     - mtl
+    - wl-pprint
   exposed-modules:
     - Control.Edit
     - Control.Lens.Edit
diff --git a/edit-lens/src/Control/DFST.lhs b/edit-lens/src/Control/DFST.lhs
index 94df533..9bd1629 100644
--- a/edit-lens/src/Control/DFST.lhs
+++ b/edit-lens/src/Control/DFST.lhs
@@ -46,7 +46,7 @@ toFST :: forall state input output. (Ord state, Ord input, Ord output) => DFST s
 toFST DFST{..} = flip execState initialFST $ forM_ (Map.toList stTransition) handleTransition
   where
     initialFST = FST
-      { stInitial = (stInitial, Nothing)
+      { stInitial = Set.singleton (stInitial, Nothing)
       , stTransition = Map.empty
       , stAccept = Set.map (, Nothing) stAccept
       }
diff --git a/edit-lens/src/Control/FST.lhs b/edit-lens/src/Control/FST.lhs
index d3c8ca9..d37072f 100644
--- a/edit-lens/src/Control/FST.lhs
+++ b/edit-lens/src/Control/FST.lhs
@@ -1,24 +1,134 @@
 \begin{code}
+{-# LANGUAGE ScopedTypeVariables
+#-}
 
 {-|
 Description: Finite state transducers with epsilon-transitions
 -}
 module Control.FST
   ( FST(..)
+  -- * Constructing FSTs
+  , wordFST
+  -- * Operations on FSTs
+  , productFST, invertFST, restrictFST
+  -- * Debugging Utilities
+  , liveFST
   ) where
 
-import Data.Map.Strict (Map)
+import Data.Map.Strict (Map, (!?))
 import qualified Data.Map.Strict as Map
 
 import Data.Set (Set)
+import qualified Data.Set as Set
 
 import Data.Sequence (Seq)
+import qualified Data.Sequence as Seq
+
+import Data.Maybe (mapMaybe, fromMaybe)
+
+import Numeric.Natural
 
 import Control.Lens.TH
 
+import Control.Monad.State.Strict
+
+import Text.PrettyPrint.Leijen (Pretty(..))
+import qualified Text.PrettyPrint.Leijen as PP
+
+import Debug.Trace
+
 data FST state input output = FST
-  { stInitial :: state
+  { stInitial :: Set state
   , stTransition :: Map (state, Maybe input) (Set (state, Maybe output))
   , stAccept :: Set state
+  } deriving (Show, Read)
+
+instance (Show state, Show input, Show output) => Pretty (FST state input output) where
+  pretty FST{..} = PP.vsep
+    [ PP.text "Initial states:" PP.</> PP.hang 2 (PP.list . map (PP.text . show) $ Set.toAscList stInitial)
+    , PP.text "State transitions:" PP.<$> PP.indent 2 (PP.vsep
+        [        PP.text (show st)
+          PP.<+> (PP.text "-" PP.<> PP.tupled [label inS, label outS] PP.<> PP.text "→")
+          PP.<+> PP.text (show st')
+        | ((st, inS), to) <- Map.toList stTransition
+        , (st', outS) <- Set.toAscList to
+        ])
+    , PP.text "Accepting states:" PP.</> PP.hang 2 (PP.list . map (PP.text . show) $ Set.toAscList stAccept)
+    ]
+    where
+      label :: Show a => Maybe a -> PP.Doc
+      label = maybe (PP.text "ɛ") (PP.text . show)
+
+wordFST :: forall input output. Seq output -> FST Natural input output
+-- ^ @wordFST str@ is the minimal FST generating @str@ as output when given no input
+wordFST outs = FST
+  { stInitial    = Set.singleton 0
+  , stAccept     = Set.singleton l
+  , stTransition = Map.fromSet next states
   }
+  where
+    l :: Natural
+    l = fromIntegral $ Seq.length outs
+    states :: Set (Natural, Maybe input)
+    states = Set.fromDistinctAscList [ (n, Nothing) | n <- [0..pred l] ]
+    next :: (Natural, Maybe input) -> Set (Natural, Maybe output)
+    next (i, _) = Set.singleton (succ i, Just . Seq.index outs $ fromIntegral i)
+
+productFST :: forall state1 state2 input output. (Ord state1, Ord state2, Ord input, Ord output) => FST state1 input output -> FST state2 input output -> FST (state1, state2) input output
+-- ^ Cartesian product on states, logical conjunction on transitions and state-properties (initial and accept)
+--
+-- This is most intuitive when thought of as the component-wise product of weighted FSTs with weights in the boolean semiring.
+productFST fst1 fst2 = FST
+  { stInitial    = stInitial fst1 `setProduct` stInitial fst2
+  , stAccept     = stAccept fst1 `setProduct` stAccept fst2
+  , stTransition = Map.fromSet transitions . Set.fromList . mapMaybe filterTransition . Set.toAscList $ Map.keysSet (stTransition fst1) `setProduct` Map.keysSet (stTransition fst2)
+  }
+  where
+    setProduct :: forall a b. Set a -> Set b -> Set (a, b)
+    setProduct as bs = Set.fromDistinctAscList $ (,) <$> Set.toAscList as <*> Set.toAscList bs
+    filterTransition :: forall label. Eq label => ((state1, Maybe label), (state2, Maybe label)) -> Maybe ((state1, state2), Maybe label)
+    filterTransition ((st1, Nothing ), (st2, in2     )) = Just ((st1, st2), in2)
+    filterTransition ((st1, in1     ), (st2, Nothing )) = Just ((st1, st2), in1)
+    filterTransition ((st1, Just in1), (st2, Just in2))
+      | in1 == in2 = Just ((st1, st2), Just in1)
+      | otherwise  = Nothing
+    transitions :: ((state1, state2), Maybe input) -> Set ((state1, state2), Maybe output)
+    transitions ((st1, st2), inS) = Set.fromList . mapMaybe filterTransition . Set.toAscList $ out1 `setProduct` out2
+      where
+        out1 = (fromMaybe Set.empty $ stTransition fst1 !? (st1, inS)) `Set.union` (fromMaybe Set.empty $ stTransition fst1 !? (st1, Nothing))
+        out2 = (fromMaybe Set.empty $ stTransition fst2 !? (st2, inS)) `Set.union` (fromMaybe Set.empty $ stTransition fst2 !? (st2, Nothing))
+
+restrictFST :: forall state input output. (Ord state, Ord input, Ord output) => Set state -> FST state input output -> FST state input output
+restrictFST sts FST{..} = FST
+  { stInitial    = stInitial `Set.intersection` sts
+  , stAccept     = stAccept  `Set.intersection` sts
+  , stTransition = Map.mapMaybeWithKey restrictTransition stTransition
+  }
+  where
+    restrictTransition :: (state, Maybe input) -> Set (state, Maybe output) -> Maybe (Set (state, Maybe output))
+    restrictTransition (st, _) tos = tos' <$ guard (st `Set.member` sts)
+      where
+        tos' = Set.filter (\(st', _) -> st' `Set.member` sts) tos
+      
+
+liveFST :: forall state input output. (Ord state, Ord input, Ord output, Show state) => FST state input output -> Set state
+-- ^ Compute the set of "live" states (with no particular complexity)
+--
+-- A state is "live" iff there is a path from it to an accepting state and a path from an initial state to it
+liveFST fst@FST{..} = flip execState Set.empty $ mapM_ (depthSearch Set.empty) stInitial
+  where
+    stTransition' :: Map state (Set state)
+    stTransition' = Map.map (Set.map $ \(st, _) -> st) $ Map.mapKeysWith Set.union (\(st, _) -> st) stTransition
+    depthSearch :: Set state -> state -> State (Set state) ()
+    depthSearch acc curr = do
+      let acc' = Set.insert curr acc
+          next = fromMaybe Set.empty $ stTransition' !? curr
+      alreadyLive <- get
+      when (not . Set.null $ Set.insert curr next `Set.intersection` Set.union stAccept alreadyLive) $
+        modify $ Set.union acc'
+      mapM_ (depthSearch acc') $ next `Set.difference` alreadyLive
+
+
+invertFST :: FST state input output -> Seq output -> Set (Seq input)
+invertFST = undefined
 \end{code}