nfa.hs

-- Like nfa.py
-- E.g.: (lit 'a' . many (lit 'b' `alt` lit 'c') . lit 'd') `matches` "abbcd"
-- Thanks to Mike Vanier for some Haskelly tips.

import           Data.Array  (Array)
import qualified Data.Array  as A
import           Data.List   (nub)
import           Data.Monoid ((<>))
import           Test.Hspec  hiding (after)

spec = hspec $ do
  describe "many" $ do
    it "should match many characters" $ do
      many (lit 'a') `shouldMatch` "aaa"
    it "should not loop forever" $ do
      many (lit 'a' `alt` lit 'a') `shouldNotMatch` "aaacf"

    it "should work on other things" $ do
      many (lit 'a' `alt` lit 'b') `shouldMatch` "abababa"
      (lit 'a' . many (lit 'b' `alt` lit 'c') . lit 'd') `shouldMatch` "abcbd"

shouldMatch re string = (re `matches` string) `shouldBe` True
shouldNotMatch re string = (re `matches` string) `shouldBe` False

type E = NFA -> NFA

type NFA  = ([Int], [NFAState])

data NFAState = Accept | Expect Char [Int]
  deriving Eq

-- how do we get a start state here?
matches :: E -> [Char] -> Bool
matches re cs =
  let (starts, states) = re ([0], [Accept])
      ends = foldl step starts cs
      step :: [Int] -> Char -> [Int]
      step starts' c = nub $ concat $ map (after c states) starts'
  in any (0 ==) ends

after :: Char -> [NFAState] -> Int -> [Int]
after c states i =
  case states !! i of
    Expect c' xs -> if c' == c then xs else []
    Accept       -> []

lit :: Char -> E
lit c (starts, states) = ([length states], states <> [Expect c starts])

alt :: E -> E -> E
alt re1 re2 nfa@(starts, states) =
  let (starts1, states1) = re1 nfa
      (starts2, states2) = re2 (starts, states1) in
  (starts1 <> starts2, states2)

-- definitely magic
many :: E -> E
many re nfa@(starts, states) =
  let (loopStarts, loopStates) = re (resultStarts, states)
      resultStarts = starts <> loopStarts in
   (resultStarts, loopStates)