The CYK algorithm in Haskell

CYK.hs

module CYK
       where

import Control.Monad
import Data.Array.IArray
import Data.Array.MArray
import Data.Array.ST
import Data.Maybe
import qualified Data.Map as M

-- General helpers.

enumerate :: [a] -> [(Int, a)]
enumerate = zip [1..]

-- Grammar ADT definition.
type Symbol   = Char
type RuleName = String

data CNFRule = TerminalRule RuleName Symbol
             | NonTerminalRule RuleName [(RuleName, RuleName)] Bool
             deriving (Eq, Show)

type CNFGrammar = [CNFRule]

-- Helpers for constructing the grammar.

ruleTerminal :: RuleName -> Symbol -> CNFRule
ruleTerminal name prod = TerminalRule name prod

ruleNonTerminal :: RuleName -> [(RuleName, RuleName)] -> CNFRule
ruleNonTerminal name prods = NonTerminalRule name prods False

ruleStart :: RuleName -> [(RuleName, RuleName)] -> CNFRule
ruleStart name prods = NonTerminalRule name prods True

-- Helper functions.

ruleName :: CNFRule -> RuleName
ruleName (TerminalRule name _)      = name
ruleName (NonTerminalRule name _ _) = name

isTerminalRule :: CNFRule -> Bool
isTerminalRule (TerminalRule _ _) = True
isTerminalRule _                  = False

isNonTerminalRule :: CNFRule -> Bool
isNonTerminalRule (NonTerminalRule _ _ _) = True
isNonTerminalRule _                       = False

isStartRule :: CNFRule -> Bool
isStartRule (NonTerminalRule _ _ b) = b
isStartRule _                       = False

terminalRules :: CNFGrammar -> [(Int, CNFRule)]
terminalRules = filter (isTerminalRule . snd) . enumerate

nonTerminalRules :: CNFGrammar -> [(Int, CNFRule)]
nonTerminalRules = filter (isNonTerminalRule . snd) . enumerate

startRules :: CNFGrammar -> [(Int, CNFRule)]
startRules = filter (isStartRule . snd) . enumerate

terminalRuleProduces :: CNFRule -> Symbol -> Bool
terminalRuleProduces (TerminalRule _ s) s' = (s == s')
terminalRuleProduces _                  _  = error "Terminal rule expected!"

lookupIndices :: (M.Map RuleName Int) -> CNFRule -> [(Int, Int)]
lookupIndices mIdx (NonTerminalRule _ prods _)
  = let lkup k = fromJust $ M.lookup k mIdx
    in [(lkup b, lkup c) | (b,c) <- prods]
lookupIndices _ _
  = error "Non-terminal rule expected!"


-- The algorithm itself
cykAlgorithm :: CNFGrammar -> [Symbol] -> Bool
cykAlgorithm grammar input = or [arr ! (1,n,x) | x <- startIndices]
  where
    n = length input
    r = length grammar
    idxMap = M.fromList (zip (map ruleName grammar) [1..])
    startIndices = map fst . startRules $ grammar

    arr = runSTUArray $ do
      marr <- newArray ((1,1,1),(n,n,r)) False

      forM_ (enumerate input) $ \(i, ci) ->
        forM_ (terminalRules grammar) $ \(j, rule) ->
          when (terminalRuleProduces rule ci) $
            writeArray marr (i,1,j) True

      forM_ [2..n] $ \i ->
        forM_ [1..(n-i+1)] $ \j ->
          forM_ [1..(i-1)] $ \k ->
            forM_ (nonTerminalRules grammar) $ \(a, rule) ->
              forM_ (lookupIndices idxMap rule) $ \(b,c) -> do
                e0 <- readArray marr (j,k,b)
                e1 <- readArray marr (j+k,i-k,c)
                when (e0 && e1) $
                  writeArray marr (j,i,a) True
      return marr

-- Example input.

-- S -> SS | LH | LR
-- H -> SR
-- L -> '('
-- R -> ')'
exampleGrammar :: CNFGrammar
exampleGrammar = [ ruleStart "S" [ ("S","S"), ("L","H"), ("L","R")]
                 , ruleNonTerminal "H" [("S","R")]
                 , ruleTerminal "L" '('
                 , ruleTerminal "R" ')'
                 ]

exampleValidInput :: [Symbol]
exampleValidInput = "((((()))))"

exampleInvalidInput :: [Symbol]
exampleInvalidInput = "(()"

-- Program entry point.
main :: IO ()
main = do let validResult = cykAlgorithm exampleGrammar exampleValidInput
          putStrLn $ "Result for the valid input: " ++ (show validResult)
          let invalidResult = cykAlgorithm exampleGrammar exampleInvalidInput
          putStrLn $ "Result for the invalid input: " ++ (show invalidResult)