{-# LANGUAGE ScopedTypeVariables
#-}

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

import Data.Map.Strict (Map, (!?))
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)
    -- ^ 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@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