{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{-|
Module      : Framework.Grammar.Parser
Description : CCG parsing, CKY-style (ish).
Copyright   : (c) Julian Grove and Aaron Steven White, 2025
License     : MIT
Maintainer  : julian.grove@gmail.com

CKY-style parsing for CCG.
-}

module Framework.Grammar.Parser ( forParse
                                , interpret
                                , interpretations
                                , ParseAs(..)
                                ) where

import           Control.Monad
import           Control.Monad.State
import           Data.List
import qualified Data.Map              as Map
import           Data.Maybe                           (maybeToList)
import           Framework.Grammar.CCG
import           Framework.Lambda
import           Prelude               as Prel hiding (Word)

--------------------------------------------------------------------------------
-- * Interface-agnostic CCG parsing

-- | Chart indices are expressions, and chart elements are sets of parses
-- (@[m]@).
type Chart m = StateT (Map.Map Expr [m]) [] ()

-- | CKY-style CCG expression interpreter, but with parses cached by substring
-- identity instead of by span, for stronger memoization.
interpret :: forall p m. Interpretation p m => Int -> Expr -> Chart m
interpret :: forall (p :: Project) m.
Interpretation p m =>
Int -> Expr -> Chart m
interpret Int
lex = \case
  -- | Cache each lexical item (but only if needed).
  [[Char]
w] -> Expr -> Chart m -> Chart m
cacheIfNeeded [[Char]
w] (Chart m -> Chart m) -> Chart m -> Chart m
forall a b. (a -> b) -> a -> b
$ do
    Map Expr [m]
m <- StateT (Map Expr [m]) [] (Map Expr [m])
forall s (m :: * -> *). MonadState s m => m s
get
    let lexSet :: [m]
lexSet = (forall (p :: Project) m. Interpretation p m => [Lexicon m]
lexica @p [Lexicon m] -> Int -> Lexicon m
forall a. HasCallStack => [a] -> Int -> a
!! Int
lex) [Char]
w
    Map Expr [m] -> Chart m
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Expr -> [m] -> Map Expr [m] -> Map Expr [m]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert [[Char]
w] [m]
lexSet Map Expr [m]
m)
  -- | Cache a complex expression (but only if needed).
  Expr
e   -> Expr -> Chart m -> Chart m
cacheIfNeeded Expr
e (Chart m -> Chart m) -> Chart m -> Chart m
forall a b. (a -> b) -> a -> b
$ do
    Int
n <- [Int] -> StateT (Map Expr [m]) [] Int
forall (m :: * -> *) a. Monad m => m a -> StateT (Map Expr [m]) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift [Int
1 .. Expr -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    let (Expr
e1, Expr
e2) = Int -> Expr -> (Expr, Expr)
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n Expr
e
    forall (p :: Project) m.
Interpretation p m =>
Int -> Expr -> Chart m
interpret @p Int
lex Expr
e1
    forall (p :: Project) m.
Interpretation p m =>
Int -> Expr -> Chart m
interpret @p Int
lex Expr
e2
    Map Expr [m]
m <- StateT (Map Expr [m]) [] (Map Expr [m])
forall s (m :: * -> *). MonadState s m => m s
get
    let i1 :: [m]
i1    = [[m]] -> [m]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe [m] -> [[m]]
forall a. Maybe a -> [a]
maybeToList (Expr -> Map Expr [m] -> Maybe [m]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr
e1 Map Expr [m]
m))
        i2 :: [m]
i2    = [[m]] -> [m]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe [m] -> [[m]]
forall a. Maybe a -> [a]
maybeToList (Expr -> Map Expr [m] -> Maybe [m]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr
e2 Map Expr [m]
m))
        i3tmp :: [m]
i3tmp = do m
i1' <- [m]
i1
                   m
i2' <- [m]
i2
                   forall (p :: Project) m. Interpretation p m => m -> m -> [m]
combineR @p m
i1' m
i2' [m] -> [m] -> [m]
forall a. [a] -> [a] -> [a]
++ forall (p :: Project) m. Interpretation p m => m -> m -> [m]
combineL @p m
i1' m
i2'
        i3 :: [m]
i3    = [[m]] -> [m]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe [m] -> [[m]]
forall a. Maybe a -> [a]
maybeToList (Expr -> Map Expr [m] -> Maybe [m]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr
e Map Expr [m]
m)) [m] -> [m] -> [m]
forall a. [a] -> [a] -> [a]
++ [m]
i3tmp 
    Map Expr [m] -> Chart m
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Expr -> [m] -> Map Expr [m] -> Map Expr [m]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Expr
e [m]
i3 Map Expr [m]
m)
  where cacheIfNeeded :: Expr -> Chart m -> Chart m
        cacheIfNeeded :: Expr -> Chart m -> Chart m
cacheIfNeeded Expr
e Chart m
chart = do Map Expr [m]
m <- StateT (Map Expr [m]) [] (Map Expr [m])
forall s (m :: * -> *). MonadState s m => m s
get
                                   case Expr -> Map Expr [m] -> Maybe [m]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr
e Map Expr [m]
m of
                                     Just [m]
_  -> () -> Chart m
forall a. a -> StateT (Map Expr [m]) [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                                     Maybe [m]
Nothing -> Chart m
chart

newtype ParseAs a = ParseAs { forall a. ParseAs a -> [a]
getList :: [a] } deriving (forall a b. (a -> b) -> ParseAs a -> ParseAs b)
-> (forall a b. a -> ParseAs b -> ParseAs a) -> Functor ParseAs
forall a b. a -> ParseAs b -> ParseAs a
forall a b. (a -> b) -> ParseAs a -> ParseAs b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> ParseAs a -> ParseAs b
fmap :: forall a b. (a -> b) -> ParseAs a -> ParseAs b
$c<$ :: forall a b. a -> ParseAs b -> ParseAs a
<$ :: forall a b. a -> ParseAs b -> ParseAs a
Functor

forParse :: ([a] -> [b]) -> ParseAs a -> ParseAs b
forParse :: forall a b. ([a] -> [b]) -> ParseAs a -> ParseAs b
forParse [a] -> [b]
f = [b] -> ParseAs b
forall a. [a] -> ParseAs a
ParseAs ([b] -> ParseAs b) -> (ParseAs a -> [b]) -> ParseAs a -> ParseAs b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [b]
f ([a] -> [b]) -> (ParseAs a -> [a]) -> ParseAs a -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseAs a -> [a]
forall a. ParseAs a -> [a]
getList

instance Show a => Show (ParseAs a) where
  show :: ParseAs a -> [Char]
show (ParseAs []) = [Char]
""
  show (ParseAs (a
a:[a]
as)) = ParseAs a -> [Char]
forall a. Show a => a -> [Char]
show ([a] -> ParseAs a
forall a. [a] -> ParseAs a
ParseAs [a]
as) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"\n" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++
                          Int -> [Char]
forall a. Show a => a -> [Char]
show ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as)) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
". " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> [Char]
forall a. Show a => a -> [Char]
show a
a 

interpretations :: forall p m. (Interpretation p m, Eq m)
                => Expr -> Int -> ParseAs m
interpretations :: forall (p :: Project) m.
(Interpretation p m, Eq m) =>
Expr -> Int -> ParseAs m
interpretations Expr
e Int
lex = [m] -> ParseAs m
forall a. [a] -> ParseAs a
ParseAs ([m] -> [m]
forall a. Eq a => [a] -> [a]
nub ([m] -> [m]) -> [m] -> [m]
forall a b. (a -> b) -> a -> b
$ [[m]] -> [m]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([[m]] -> [m]) -> [[m]] -> [m]
forall a b. (a -> b) -> a -> b
$
                                    (Map Expr [m] -> [m]) -> [Map Expr [m]] -> [[m]]
forall a b. (a -> b) -> [a] -> [b]
map ([[m]] -> [m]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([[m]] -> [m]) -> (Map Expr [m] -> [[m]]) -> Map Expr [m] -> [m]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe [m] -> [[m]]
forall a. Maybe a -> [a]
maybeToList (Maybe [m] -> [[m]])
-> (Map Expr [m] -> Maybe [m]) -> Map Expr [m] -> [[m]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Map Expr [m] -> Maybe [m]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Expr
e)
                                    (StateT (Map Expr [m]) [] () -> Map Expr [m] -> [Map Expr [m]]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (forall (p :: Project) m.
Interpretation p m =>
Int -> Expr -> Chart m
interpret @p Int
lex Expr
e) Map Expr [m]
forall k a. Map k a
Map.empty)
                                   )