Filling up for instructional automata paper in Haskell

fa.hs

{- ported from and/or filled for
  "Regular Expressions and Automata using Haskell, Simon Thompson (Jan 2000)",
   by @nfunato on 2012/04/07 -}

{- this code is provided "as-is". -}

import Data.List (foldl', nub, sort, sortBy, groupBy, elemIndex, findIndex)
import Data.Set  (Set, empty, size, singleton, fromList, toList, toAscList,
                  elems, member, findMin, insert, union, intersection,(\\))
import qualified Data.List as L -- L.map
import qualified Data.Set  as S -- S.map
import Data.Maybe (fromJust)
import Control.Exception (assert)
import Data.Ord (comparing)
import Data.Function (fix, on)

-- ===================================================================
-- finite automata
--

type Symbol = Char

data Transition a =
   Epsilon a a                  -- Epsilon  fromState       toState
 | Labelled a Symbol a          -- Labelled fromState label toState
 deriving (Eq, Ord, Show)

data Fa a = Fa {                -- Finite Automaton
   states :: Set a
 , transitions :: Set (Transition a)
 , start :: a
 , goals :: Set a
 } deriving (Eq, Show)

makeFa :: Ord a => ([a], [Transition a], a, [a]) -> Fa a
makeFa (states, transitions, start, goals) =
  assert checkArgs (Fa stSet trSet start goalSet)
    where
      stSet   = fromList states
      trSet   = fromList transitions
      goalSet = fromList goals
      checkArgs =
        let memSt st = st `elem` states
            validTr (Epsilon  f   t) = memSt f && memSt t
            validTr (Labelled f _ t) = memSt f && memSt t
        -- also should check isolated state here
        in memSt start
           && all memSt goals
           && all validTr transitions

accept :: Ord a => Fa a -> [Symbol] -> Bool
accept fa syms = any (flip member (goals fa)) states
  where states = toList $ transit fa syms

transit :: Ord a => Fa a -> [Symbol] -> Set a
transit fa syms = foldl' (transitOne fa) startStates syms
  where startStates = eClosure fa (singleton (start fa))

transitOne :: Ord a => Fa a -> Set a -> Symbol -> Set a
transitOne fa states = (eClosure fa) . transitOneLabel fa states

transitOneLabel :: Ord a => Fa a -> Set a -> Symbol -> Set a
transitOneLabel fa argStates argSym =
  fromList [ toSt | argSt <- elems argStates,
                    Labelled fromSt sym toSt <- elems (transitions fa),
                    argSt == fromSt, argSym == sym ]

transitOneEpsilon :: Ord a => Fa a -> Set a -> Set a
transitOneEpsilon fa argStates =
  fromList [ toSt | argSt <- elems argStates,
                    Epsilon fromSt toSt <- elems (transitions fa),
                    argSt == fromSt ]

eClosure :: Ord a => Fa a -> Set a -> Set a
eClosure fa argStates = fixpoint step argStates
  where step states = states `union` transitOneEpsilon fa states

fixpoint :: Eq a => (a -> a) -> a -> a
fixpoint f x
    | x == newX  = x
    | otherwise  = fixpoint f newX
  where newX = f x

-- ===================================================================
-- automata arithmetic (just compiled, not tested at all)
--

renumSt :: Enum a => Int -> a -> a
renumSt offset state = (iterate succ state) !! offset

renumTr :: Enum a => Int -> Transition a -> Transition a
renumTr offset transition =
   case transition of
     Epsilon  f   t -> Epsilon  (renum f)   (renum t)
     Labelled f s t -> Labelled (renum f) s (renum t)
 where renum = renumSt offset

ensureSingleGoal :: Fa Int -> Fa Int
ensureSingleGoal fa@(Fa stSet trSet start goalSet) =
   if size goalSet == 1
   then fa
   else Fa newStSet newTrSet start (singleton m)
 where
   m = size stSet
   newStSet = insert m stSet
   newTrSet = trSet `union` S.map (\g -> Epsilon g m) goalSet

mOr fa1 fa2 = mOr' (ensureSingleGoal fa1) (ensureSingleGoal fa2) where
 mOr' (Fa stSet1 trSet1 start1 goalSet1) (Fa stSet2 trSet2 start2 goalSet2) =
   Fa newStSet newTrSet 0 (singleton (m1+m2+1))
     where
       goal1 = (elems goalSet1) !! 0
       goal2 = (elems goalSet2) !! 0
       m1 = size stSet1
       m2 = size stSet2
       stSet1' = S.map (renumSt 1) stSet1
       stSet2' = S.map (renumSt (m1+1)) stSet2
       trSet1' = S.map (renumTr 1) trSet1
       trSet2' = S.map (renumTr (m1+1)) trSet2
       newStSet = stSet1' `union` stSet2' `union` fromList [0, (m1+m2+1)]
       newTrSet = trSet1' `union` trSet2'
                          `union` fromList [ Epsilon 0  (start1+1),
                                             Epsilon 0  (start2+m1+1),
                                             Epsilon (goal1+1) (m1+m2+1),
                                             Epsilon (goal2+m1+1)(m1+m2+1)]

mThen fa1 fa2 = mThen' (ensureSingleGoal fa1) fa2 where
 mThen' (Fa stSet1 trSet1 start1 goalSet1) (Fa stSet2 trSet2 start2 goalSet2) =
   Fa newStSet newTrSet start1 goalSet2
     where
       goal1 = (elems goalSet1) !! 0
       m1 = size stSet1
       stSet2' = S.map (renumSt m1) stSet2
       trSet2' = S.map (renumTr m1) trSet2
       newStSet = stSet1 `union` stSet2'
       newTrSet = trSet1 `union` trSet2'
                         `union` fromList [ Epsilon goal1 (start2+m1) ]

mStar fa = mStar' (ensureSingleGoal fa) where
 mStar' (Fa stSet trSet start goalSet) =
   Fa newStSet newTrSet 0 (singleton (m+1))
     where
       goal = (elems goalSet) !! 0
       m = size stSet
       stSet' = S.map (renumSt 1) stSet
       trSet' = S.map (renumTr 1) (insert (Epsilon goal start) trSet)
       newStSet = stSet' `union` fromList [0,m+1]
       newTrSet = trSet' `union` fromList [ Epsilon 0 (m+1),
                                            Epsilon 0 (start+1),
                                            Epsilon (goal+1) (m+1) ]

-- ===================================================================
-- conversion from nfa to dfa (just tested for a few examples)
--

makeDfa :: Ord a => Fa a -> Fa Int
makeDfa = renumFa . makeDfa'

renumFa :: Ord a => Fa a -> Fa Int
renumFa (Fa stsSet trsSet starts goalsSet) =
  Fa stSet trSet start goalSet
    where -- we might want to memoize the result of convSt ...
      stSet   = fromList [0 .. stSize-1]
      trSet   = S.map convTr trsSet
      start   = convSt starts
      goalSet = S.map convSt goalsSet
      stList  = toAscList stsSet
      stSize  = length stList
      convSt st = fromJust $ elemIndex st stList
      convTr (Labelled f s t) = Labelled (convSt f) s (convSt t)
--    convTr (Epsilon  f   t) = Epsilon  (convSt f)   (convSt t)

-- maybe we should incorporate 'alphabet' into a member of record 'Fa'.
alphabet :: Fa a -> [Symbol]
alphabet fa = sort $ nub [ sym | Labelled _ sym _ <- elems (transitions fa) ]

makeDfa' :: Ord a => Fa a -> Fa (Set a)
makeDfa' nfa = fixpoint (addStep nfa (alphabet nfa)) startDfa
  where
    startDfa      = Fa (singleton startMetaSt) empty startMetaSt goalMetaSts
    startMetaSt   = eClosure nfa (singleton (start nfa))
    goalMetaSts
      | (goals nfa) `intersection` startMetaSt == empty = empty
      | otherwise = singleton startMetaSt

addStep :: Ord a => Fa a -> [Symbol] -> Fa (Set a) -> Fa (Set a)
addStep nfa syms dfa =
  foldl' (\dfa' dfaState' -> addTrs nfa dfaState' dfa' syms) dfa dfaStates
    where dfaStates = toList (states dfa)

addTrs :: Ord a => Fa a -> Set a -> Fa (Set a) -> [Symbol] -> Fa (Set a)
addTrs nfa nfaStates dfa syms = foldl' (addTr nfa nfaStates) dfa syms

addTr :: Ord a => Fa a -> Set a -> Fa (Set a) -> Symbol -> Fa (Set a)
addTr nfa nfaStates (Fa stSet trSet start goalSet) sym =
 Fa stSet' trSet' start goalSet'
   where
     newStates     = transitOne nfa nfaStates sym
     stSet'        = stSet `union` singleton newStates
     trSet'        = trSet `union` singleton (Labelled nfaStates sym newStates)
     goalSet'
       | empty /= (goals nfa) `intersection` newStates
                   = goalSet `union` singleton newStates
       | otherwise = goalSet

-- ===================================================================
-- minimizing dfa (just tested for a few examples)
--

-- maybe we can unify the body of 'minimizeDfa' into 'renumFa'.
minimizeDfa :: Ord a => Fa a -> Fa Int
minimizeDfa dfa@(Fa stSet trSet start goalSet)= Fa stSet' trSet' start' goalSet'
  where -- we might want to memoize the result of convSt ...
    stSet'    = fromList [0 .. (length partition)-1]
    trSet'    = S.map convTr trSet
    start'    = convSt start
    goalSet'  = S.map convSt goalSet
    partition = sortBy (comparing findMin) $ partitionDfaStates dfa
    convSt st = fromJust $ findIndex (st `member`) partition
    convTr (Labelled f s t) = Labelled (convSt f) s (convSt t)

partitionDfaStates :: Ord a => Fa a -> [Set a]
partitionDfaStates (Fa stSet trSet _ goalSet) = fixpoint step initialPartition
  where
    initialPartition = [goalSet, stSet\\goalSet]
    step partition = refinePartitionBy groupId partition
      where 
        groupId state = fromList [ (sym, inNthPartition toSt) |
                                   Labelled frSt sym toSt <- elems trSet,
                                   frSt == state ]
        inNthPartition state = fromJust $ findIndex (state `member`) partition

refinePartitionBy :: (Ord a, Ord b) => (a -> b) -> [Set a] -> [Set a]
refinePartitionBy groupId partition = concatMap refineElemOf partition
  where refineElemOf =
          L.map (fromList . (L.map fst))
            . groupBy ((==) `on` snd) . sortBy (comparing snd)
              . L.map (\x -> (x, groupId x)) . toList

-- ===================================================================
-- test data
--

nfa1    = makeFa ([ 0..3 ],
                  [ Labelled 0 'a' 0,
                    Labelled 0 'a' 1,
                    Labelled 0 'b' 0,
                    Labelled 1 'b' 2,
                    Labelled 2 'b' 3 ],
                  0,
                  [3])

nfa2    = makeFa ([ 0..5 ],
                  [ Labelled 0 'a' 1,
                    Labelled 1 'b' 2,
                    Labelled 0 'a' 3,
                    Labelled 3 'b' 4,
                    Epsilon  3  4,
                    Labelled 4 'b' 5 ],
                  0,
                  [2,5])

nfaEx11 = makeFa ([ 0 .. 8],
                  [ Epsilon  0  1,
                    Epsilon  1  2,
                    Labelled 2 'a' 3,
                    Epsilon  3  6,
                    Epsilon  6  1,
                    Epsilon  1  4,
                    Labelled 4 'b' 5,
                    Epsilon  5  6,
                    Epsilon  6  7,
                    Labelled 7 'a' 8 ],
                  0,
                  [8])

-- aka machineP
mP = nfaEx11

-- eof