{-# LANGUAGE RecordWildCards
           , PatternGuards
           , ScopedTypeVariables
#-}

{-|
Description: Deterministic finite state transducers
-}
module Data.String.DFST
  ( DFST(..)
  , runDFST
  ) where

import Data.Map.Strict (Map, (!?))
import qualified Data.Map.Strict as Map

import Data.Set (Set)
import qualified Data.Set as Set

data DFST state = DFST
  { stInitial :: state
  , stTransition :: Map (state, Char) (state, Maybe Char)
    -- ^ All @(s, c)@-combinations not mapped are assumed to map to @(s, Nothing)@
  , stAccept :: Set state
  }
  
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