NDFSM toy implementation

Ndfsm.hs

module Ndfsm
    ( State (..)
    , Transitions
    , Ndfsm

    , fromTransition
    , fromTransitions
    , fromAccepting
    , fromStart
    , lookupTransitions
    , eps

    , TransitionsM
    , (-->)
    , epsTransitions

    , NdfsmM
    , newState
    , transitions
    , accepting
    , start

    , ndsfm

    , simulate
    ) where

import Control.Monad (mplus)
import Control.Monad.Writer (Writer, WriterT, tell, execWriter, execWriterT)
import Data.Map (Map)
import Data.Maybe (fromMaybe)
import Data.Monoid (Monoid, mempty, mappend, mconcat)
import Data.Set (Set)
import qualified Control.Monad.State as MS (State, get, put, evalState)
import qualified Data.Map as M
import qualified Data.Set as S

--------------------------------------------------------------------------------
-- Core types                                                                 --
--------------------------------------------------------------------------------

-- | Identifies a state
--
newtype State = State {unState :: Int}
              deriving (Eq, Ord, Show)

-- | A transition table for a state
--
newtype Transitions a = Transitions {unTransitions :: Map (Maybe a) (Set State)}
                      deriving (Eq, Ord, Show)

instance Ord a => Monoid (Transitions a) where
    mempty = Transitions M.empty
    mappend x y = Transitions $
        M.unionWith S.union (unTransitions x) (unTransitions y)

-- | Non-deterministic finite state machine
--
data Ndfsm a = Ndfsm
    { ndfsmTransitions :: Map State (Transitions a)
    , ndfsmAccepting   :: Set State
    , ndfsmStart       :: Maybe State
    } deriving (Show)

instance Ord a => Monoid (Ndfsm a) where
    mempty = Ndfsm M.empty S.empty Nothing
    mappend x y = Ndfsm
        (ndfsmTransitions x `mappend` ndfsmTransitions y)
        (ndfsmAccepting x `mappend` ndfsmAccepting y)
        (ndfsmStart x `mplus` ndfsmStart y)

--------------------------------------------------------------------------------
-- Core operations                                                            --
--------------------------------------------------------------------------------

fromTransition :: Ord a => Maybe a -> State -> Transitions a
fromTransition x = Transitions . M.singleton x . S.singleton

fromTransitions :: Ord a => State -> Transitions a -> Ndfsm a
fromTransitions s t = Ndfsm (M.singleton s t) S.empty Nothing

fromAccepting :: Ord a => Set State -> Ndfsm a
fromAccepting a = mempty {ndfsmAccepting = a}

fromStart :: Ord a => State -> Ndfsm a
fromStart s = mempty {ndfsmStart = Just s}

lookupTransitions :: Ord a => State -> Maybe a -> Ndfsm a -> Set State
lookupTransitions s x m = fromMaybe S.empty $ do
    t <- M.lookup s $ ndfsmTransitions m
    M.lookup x $ unTransitions t

eps :: Ord a => State -> Ndfsm a -> Set State
eps s m = eps' (S.singleton s) [s]
  where
    eps' r [] = r
    eps' r n  =
        let n' = S.unions $ map (\x -> lookupTransitions x Nothing m) n
        in eps' (r `S.union` n') (S.toList $ n' `S.difference` r)

--------------------------------------------------------------------------------
-- Construction DSL for transitions                                           --
--------------------------------------------------------------------------------

type TransitionsM a b = Writer (Transitions a) b

(-->) :: Ord a => a -> State -> TransitionsM a ()
(-->) x s = tell $ fromTransition (Just x) s

epsTransitions :: Ord a => [State] -> TransitionsM a ()
epsTransitions = tell . mconcat . map (fromTransition Nothing)

--------------------------------------------------------------------------------
-- Construction DSL for NDFSM's                                               --
--------------------------------------------------------------------------------

type NdfsmM a b = WriterT (Ndfsm a) (MS.State Int) b

newState :: Ord a => NdfsmM a State
newState = do
    i <- MS.get
    MS.put $ i + 1
    return $ State {unState = i}

transitions :: Ord a => State -> TransitionsM a () -> NdfsmM a ()
transitions s = tell . fromTransitions s . execWriter

accepting :: Ord a => [State] -> NdfsmM a ()
accepting = tell . fromAccepting . S.fromList

start :: Ord a => State -> NdfsmM a ()
start = tell . fromStart

ndsfm :: Ord a => NdfsmM a () -> Ndfsm a
ndsfm n = MS.evalState (execWriterT n) 0

--------------------------------------------------------------------------------
-- Simulation                                                                 --
--------------------------------------------------------------------------------

simulate :: Ord a => Ndfsm a -> [a] -> Bool
simulate m input = case ndfsmStart m of
    Nothing     -> error "No starting state specified!"
    Just start' -> simulate' (eps start' m) input
  where
    -- No more input, check that at least one state is accepting
    simulate' st [] = not $ S.null $ st `S.intersection` ndfsmAccepting m
    simulate' st (c : cs)
        | S.null st = False  -- Shortcut quit when no more states are found
        | otherwise =
            let st' = S.unions $ map (\x -> eps x m) $ S.toList
                    $ S.unions $ map (\x -> lookupTransitions x (Just c) m)
                    $ S.toList st
            in simulate' st' cs

test.hs

import Ndfsm

import Control.Monad (replicateM)

noMoreThanOneB :: Ndfsm Char
noMoreThanOneB = ndsfm $ do
    s <- newState
    t <- newState

    start s
    accepting [s, t]

    transitions s $ do
        'a' --> s
        'b' --> t

    transitions t $ do
        'a' --> t

missingLetter :: Ndfsm Char
missingLetter = ndsfm $ do
    [q0, notA, notB, notC, notD] <- replicateM 5 newState

    start q0
    accepting [notA, notB, notC, notD]

    transitions q0 $ do
        epsTransitions [notA, notB, notC, notD]

    transitions notA $ do
        'b' --> notA
        'c' --> notA
        'd' --> notA

    transitions notB $ do
        'a' --> notB
        'c' --> notB
        'd' --> notB

    transitions notC $ do
        'a' --> notC
        'b' --> notC
        'd' --> notC

    transitions notD $ do
        'a' --> notD
        'b' --> notD
        'c' --> notD