{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Predict sets of a context-free grammar.
module Data.Cfg.Internal.PredictSet
  ( Prediction
  , Predictions
  , predictSet
  , ll1InfoMap
  , isLL1
  ) where

import Data.Cfg.Augment
import Data.Cfg.Cfg (Cfg(..))
import Data.Cfg.Collect
import Data.Cfg.Internal.FirstSet (firstsOfVs)
import Data.Cfg.LookaheadSet
import qualified Data.Map.Strict as M
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import qualified Data.Set as S

-- | Returns the predict set of a production.
predictSet ::
     (Ord t)
  => (AugNT nt -> LookaheadSet t) -- ^ 'firstSet' for the grammar
  -> (AugNT nt -> LookaheadSet t) -- ^ 'followSet' for the grammar
  -> AugProduction t nt -- ^ the production
  -> LookaheadSet t
predictSet :: (AugNT nt -> LookaheadSet t)
-> (AugNT nt -> LookaheadSet t)
-> AugProduction t nt
-> LookaheadSet t
predictSet AugNT nt -> LookaheadSet t
firstSet' AugNT nt -> LookaheadSet t
followSet' (AugNT nt
hd, Vs (AugT t) (AugNT nt)
vs) =
  (AugNT nt -> LookaheadSet t)
-> Vs (AugT t) (AugNT nt) -> LookaheadSet t
forall t nt.
Ord t =>
(AugNT nt -> LookaheadSet t) -> AugVs t nt -> LookaheadSet t
firstsOfVs AugNT nt -> LookaheadSet t
firstSet' Vs (AugT t) (AugNT nt)
vs LookaheadSet t -> LookaheadSet t -> LookaheadSet t
forall a. Semigroup a => a -> a -> a
<> AugNT nt -> LookaheadSet t
followSet' AugNT nt
hd

-- | A lookahead set with the productions it predicts
type Prediction t nt = (LookaheadSet t, S.Set (AugProduction t nt))

-- | A set of 'Prediction's.  The 'LookaheadSet's of the 'Prediction's
-- will be pairwise disjoint.
type Predictions t nt = S.Set (Prediction t nt)

-- | Returns the production 'Predictions' for the grammar as a map.
ll1InfoMap ::
     forall cfg t nt. (Cfg cfg (AugT t) (AugNT nt), Ord nt, Ord t)
  => cfg (AugT t) (AugNT nt)
  -> (AugProduction t nt -> LookaheadSet t)
  -> M.Map (AugNT nt) (Predictions t nt)
ll1InfoMap :: cfg (AugT t) (AugNT nt)
-> (AugProduction t nt -> LookaheadSet t)
-> Map (AugNT nt) (Predictions t nt)
ll1InfoMap cfg (AugT t) (AugNT nt)
cfg AugProduction t nt -> LookaheadSet t
predictSet' = (AugNT nt -> Predictions t nt)
-> [AugNT nt] -> Map (AugNT nt) (Predictions t nt)
forall k v. Ord k => (k -> v) -> [k] -> Map k v
mkMap AugNT nt -> Predictions t nt
mkPredictions ([AugNT nt] -> Map (AugNT nt) (Predictions t nt))
-> [AugNT nt] -> Map (AugNT nt) (Predictions t nt)
forall a b. (a -> b) -> a -> b
$ Set (AugNT nt) -> [AugNT nt]
forall a. Set a -> [a]
S.toList (Set (AugNT nt) -> [AugNT nt]) -> Set (AugNT nt) -> [AugNT nt]
forall a b. (a -> b) -> a -> b
$ cfg (AugT t) (AugNT nt) -> Set (AugNT nt)
forall (cfg :: * -> * -> *) t nt.
Cfg cfg t nt =>
cfg t nt -> Set nt
nonterminals cfg (AugT t) (AugNT nt)
cfg
  where
    mkPredictions :: AugNT nt -> Predictions t nt
    -- Mostly reshuffling data
    mkPredictions :: AugNT nt -> Predictions t nt
mkPredictions AugNT nt
nt =
      [(LookaheadSet t, Set (AugProduction t nt))] -> Predictions t nt
forall a. Ord a => [a] -> Set a
S.fromList ([(LookaheadSet t, Set (AugProduction t nt))] -> Predictions t nt)
-> [(LookaheadSet t, Set (AugProduction t nt))] -> Predictions t nt
forall a b. (a -> b) -> a -> b
$ [([AugT t], Set (AugProduction t nt))]
-> [(LookaheadSet t, Set (AugProduction t nt))]
f ([([AugT t], Set (AugProduction t nt))]
 -> [(LookaheadSet t, Set (AugProduction t nt))])
-> [([AugT t], Set (AugProduction t nt))]
-> [(LookaheadSet t, Set (AugProduction t nt))]
forall a b. (a -> b) -> a -> b
$ [(AugT t, Set (AugProduction t nt))]
-> [([AugT t], Set (AugProduction t nt))]
forall b a. Eq b => [(a, b)] -> [([a], b)]
collectOnSecond ([(AugT t, Set (AugProduction t nt))]
 -> [([AugT t], Set (AugProduction t nt))])
-> [(AugT t, Set (AugProduction t nt))]
-> [([AugT t], Set (AugProduction t nt))]
forall a b. (a -> b) -> a -> b
$ [(AugT t, AugProduction t nt)]
-> [(AugT t, Set (AugProduction t nt))]
forall a b. (Eq a, Ord b) => [(a, b)] -> [(a, Set b)]
collectOnFirst' [(AugT t, AugProduction t nt)]
lookaheadProds
      -- Possible lookahead symbols for productions of this nonterminal
      where
        lookaheadProds :: [(AugT t, AugProduction t nt)]
        lookaheadProds :: [(AugT t, AugProduction t nt)]
lookaheadProds = do
          Vs (AugT t) (AugNT nt)
rhs <- Set (Vs (AugT t) (AugNT nt)) -> [Vs (AugT t) (AugNT nt)]
forall a. Set a -> [a]
S.toList (Set (Vs (AugT t) (AugNT nt)) -> [Vs (AugT t) (AugNT nt)])
-> Set (Vs (AugT t) (AugNT nt)) -> [Vs (AugT t) (AugNT nt)]
forall a b. (a -> b) -> a -> b
$ cfg (AugT t) (AugNT nt) -> AugNT nt -> Set (Vs (AugT t) (AugNT nt))
forall (cfg :: * -> * -> *) t nt.
Cfg cfg t nt =>
cfg t nt -> nt -> Set (Vs t nt)
productionRules cfg (AugT t) (AugNT nt)
cfg AugNT nt
nt
          let prod :: AugProduction t nt
prod = (AugNT nt
nt, Vs (AugT t) (AugNT nt)
rhs)
          AugT t
t <- Set (AugT t) -> [AugT t]
forall a. Set a -> [a]
S.toList (Set (AugT t) -> [AugT t]) -> Set (AugT t) -> [AugT t]
forall a b. (a -> b) -> a -> b
$ LookaheadSet t -> Set (AugT t)
forall t. LookaheadSet t -> Set (AugT t)
toSet (LookaheadSet t -> Set (AugT t)) -> LookaheadSet t -> Set (AugT t)
forall a b. (a -> b) -> a -> b
$ AugProduction t nt -> LookaheadSet t
predictSet' AugProduction t nt
prod
          (AugT t, AugProduction t nt) -> [(AugT t, AugProduction t nt)]
forall (m :: * -> *) a. Monad m => a -> m a
return (AugT t
t, AugProduction t nt
prod)
        f :: [([AugT t], S.Set (AugProduction t nt))]
          -> [(LookaheadSet t, S.Set (AugProduction t nt))]
        f :: [([AugT t], Set (AugProduction t nt))]
-> [(LookaheadSet t, Set (AugProduction t nt))]
f [([AugT t], Set (AugProduction t nt))]
pairs = [([AugT t] -> LookaheadSet t
forall t. Ord t => [AugT t] -> LookaheadSet t
fromList [AugT t]
la, Set (AugProduction t nt)
ps) | ([AugT t]
la, Set (AugProduction t nt)
ps) <- [([AugT t], Set (AugProduction t nt))]
pairs]
    mkMap :: Ord k => (k -> v) -> [k] -> M.Map k v
    mkMap :: (k -> v) -> [k] -> Map k v
mkMap k -> v
f [k]
ks = [(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k
k, k -> v
f k
k) | k
k <- [k]
ks]

-- | Returns true iff the predictions are unambiguous, true iff the
-- grammar is LL(1).
isLL1 :: M.Map (AugNT nt) (Predictions t nt) -> Bool
isLL1 :: Map (AugNT nt) (Predictions t nt) -> Bool
isLL1 Map (AugNT nt) (Predictions t nt)
m = (Predictions t nt -> Bool) -> [Predictions t nt] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Predictions t nt
ps -> Predictions t nt -> Int
forall a. Set a -> Int
S.size Predictions t nt
ps Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) ([Predictions t nt] -> Bool) -> [Predictions t nt] -> Bool
forall a b. (a -> b) -> a -> b
$ Map (AugNT nt) (Predictions t nt) -> [Predictions t nt]
forall k a. Map k a -> [a]
M.elems Map (AugNT nt) (Predictions t nt)
m