cky.hs

import Data.Array
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Text.Printf (printf)

data Symbol = S | N | V | P | D | NP | VP | PP
  deriving (Show, Eq, Ord)

data ParseTree
  = Leaf Symbol String
  | Unary Symbol ParseTree
  | Binary Symbol ParseTree ParseTree
  deriving (Eq, Ord, Show)

type TerminalMap = M.Map String [Symbol]
type UnaryMap = M.Map Symbol [Symbol]
type BinaryMap = M.Map (Symbol, Symbol) [Symbol]

terminals :: TerminalMap
terminals =
  M.fromList
    [ ("time", [V, N])
    , ("flies", [V, N])
    , ("like", [V, N, P])
    , ("arrow", [N])
    , ("an", [D])
    ]

unaryRules :: UnaryMap
unaryRules =
  M.fromList
    [ (VP, [S])
    , (NP, [S])
    , (V, [VP])
    , (N, [NP])
    ]

binaryRules :: BinaryMap
binaryRules =
  M.fromList
    [ ((NP, VP), [S])
    , ((V, NP), [VP])
    , ((VP, PP), [VP])
    , ((D, N), [NP])
    , ((N, N), [NP])
    , ((NP, PP), [NP])
    , ((P, NP), [PP])
    ]

rootSymbol :: ParseTree -> Symbol
rootSymbol (Leaf s _) = s
rootSymbol (Unary s _) = s
rootSymbol (Binary s _ _) = s

unaryClosure :: ParseTree -> [ParseTree]
unaryClosure root = bfs [root] S.empty
 where
  bfs [] _ = []
  bfs (t : ts) visited
    | t `S.member` visited = bfs ts visited
    | otherwise = t : bfs (parents ++ ts) (S.insert t visited)
   where
    sym = rootSymbol t
    parents = [Unary p t | p <- M.findWithDefault [] sym unaryRules]

initialize :: String -> [ParseTree]
initialize token = concatMap unaryClosure leaves
 where
  leaves = [Leaf sym token | sym <- M.findWithDefault [] token terminals]

parse :: [String] -> [ParseTree]
parse tokens = dp ! (1, n)
 where
  n = length tokens

  dp :: Array (Int, Int) [ParseTree]
  dp =
    array
      ((1, 1), (n, n))
      [((i, j), solve i j) | i <- [1 .. n], j <- [i .. n]]

  solve :: Int -> Int -> [ParseTree]
  solve i j
    | i == j = initialize (tokens !! (i - 1))
    | otherwise = concatMap unaryClosure $ do
        k <- [i .. j - 1]
        left <- dp ! (i, k)
        right <- dp ! (k + 1, j)
        let lsym = rootSymbol left
        let rsym = rootSymbol right
        sym <- M.findWithDefault [] (lsym, rsym) binaryRules
        return $ Binary sym left right

quoted :: (Show a) => a -> String
quoted = show . show

pretty :: ParseTree -> String
pretty (Leaf v keyword) = printf "tree(%s, %s)" (quoted v) (show keyword)
pretty (Unary v sub) = printf "tree(%s, %s)" (quoted v) (pretty sub)
pretty (Binary v l r) = printf "tree(%s, %s, %s)" (quoted v) (pretty l) (pretty r)

s = "time flies like an arrow"

main =
  mapM_ (putStrLn . pretty) (parse (words s))

$ runghc cyk.hs
tree("VP", tree("V", "time"), tree("NP", tree("NP", tree("N", "flies")), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow")))))
tree("S", tree("VP", tree("V", "time"), tree("NP", tree("NP", tree("N", "flies")), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow"))))))
tree("S", tree("NP", tree("N", "time")), tree("VP", tree("VP", tree("V", "flies")), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow")))))
tree("VP", tree("VP", tree("V", "time"), tree("NP", tree("N", "flies"))), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow"))))
tree("S", tree("VP", tree("VP", tree("V", "time"), tree("NP", tree("N", "flies"))), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow")))))
tree("S", tree("NP", tree("N", "time"), tree("N", "flies")), tree("VP", tree("V", "like"), tree("NP", tree("D", "an"), tree("N", "arrow"))))
tree("NP", tree("NP", tree("N", "time"), tree("N", "flies")), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow"))))
tree("S", tree("NP", tree("NP", tree("N", "time"), tree("N", "flies")), tree("PP", tree("P", "like"), tree("NP", tree("D", "an"), tree("N", "arrow")))))