-- | First sets of a context-free grammar.
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Cfg.Internal.FirstSet
  ( firstSetMap
  , firstsOfVs
  ) where

import Data.Cfg.Augment
import Data.Cfg.Cfg
import Data.Cfg.FixedPoint (fixedPoint)
import Data.Cfg.LookaheadSet hiding (unions)
import qualified Data.Cfg.LookaheadSet as LA
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe)
import Data.Monoid (Monoid(..))
import qualified Data.Set as S

-- | Returns the first set of the nonterminal for the grammar as a
-- map.
firstSetMap ::
     forall cfg t nt. (Cfg cfg (AugT t) (AugNT nt), Ord nt, Ord t)
  => cfg (AugT t) (AugNT nt)
  -> M.Map (AugNT nt) (LookaheadSet t)
firstSetMap :: cfg (AugT t) (AugNT nt) -> Map (AugNT nt) (LookaheadSet t)
firstSetMap cfg (AugT t) (AugNT nt)
cfg = (Map (AugNT nt) (LookaheadSet t)
 -> Map (AugNT nt) (LookaheadSet t))
-> Map (AugNT nt) (LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
forall a. Eq a => (a -> a) -> a -> a
fixedPoint Map (AugNT nt) (LookaheadSet t) -> Map (AugNT nt) (LookaheadSet t)
go Map (AugNT nt) (LookaheadSet t)
forall k a. Map k a
M.empty
  where
    go :: M.Map (AugNT nt) (LookaheadSet t) -> M.Map (AugNT nt) (LookaheadSet t)
    go :: Map (AugNT nt) (LookaheadSet t) -> Map (AugNT nt) (LookaheadSet t)
go Map (AugNT nt) (LookaheadSet t)
knownFirsts =
      [(AugNT nt, LookaheadSet t)] -> Map (AugNT nt) (LookaheadSet t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        [ (AugNT nt
nt, [Vs (AugT t) (AugNT nt)] -> LookaheadSet t
firstAlts [Vs (AugT t) (AugNT nt)]
rhss)
        | AugNT nt
nt <- 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
        , let rhss :: [Vs (AugT t) (AugNT nt)]
rhss = 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
        , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Vs (AugT t) (AugNT nt)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Vs (AugT t) (AugNT nt)]
rhss
        ]
      where
        firstAlts :: [Vs (AugT t) (AugNT nt)] -> LookaheadSet t
        firstAlts :: [Vs (AugT t) (AugNT nt)] -> LookaheadSet t
firstAlts = [LookaheadSet t] -> LookaheadSet t
forall t. Ord t => [LookaheadSet t] -> LookaheadSet t
LA.unions ([LookaheadSet t] -> LookaheadSet t)
-> ([Vs (AugT t) (AugNT nt)] -> [LookaheadSet t])
-> [Vs (AugT t) (AugNT nt)]
-> LookaheadSet t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vs (AugT t) (AugNT nt) -> LookaheadSet t)
-> [Vs (AugT t) (AugNT nt)] -> [LookaheadSet t]
forall a b. (a -> b) -> [a] -> [b]
map ([LookaheadSet t] -> LookaheadSet t
forall a. Monoid a => [a] -> a
mconcat ([LookaheadSet t] -> LookaheadSet t)
-> (Vs (AugT t) (AugNT nt) -> [LookaheadSet t])
-> Vs (AugT t) (AugNT nt)
-> LookaheadSet t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (V (AugT t) (AugNT nt) -> LookaheadSet t)
-> Vs (AugT t) (AugNT nt) -> [LookaheadSet t]
forall a b. (a -> b) -> [a] -> [b]
map (Map (AugNT nt) (LookaheadSet t)
-> V (AugT t) (AugNT nt) -> LookaheadSet t
forall nt t.
Ord nt =>
Map (AugNT nt) (LookaheadSet t)
-> V (AugT t) (AugNT nt) -> LookaheadSet t
firstsV Map (AugNT nt) (LookaheadSet t)
knownFirsts))

firstsV ::
     Ord nt
  => M.Map (AugNT nt) (LookaheadSet t)
  -> V (AugT t) (AugNT nt)
  -> LookaheadSet t
firstsV :: Map (AugNT nt) (LookaheadSet t)
-> V (AugT t) (AugNT nt) -> LookaheadSet t
firstsV Map (AugNT nt) (LookaheadSet t)
_ (T AugT t
t) = AugT t -> LookaheadSet t
forall t. AugT t -> LookaheadSet t
LA.singleton AugT t
t
firstsV Map (AugNT nt) (LookaheadSet t)
fs (NT AugNT nt
nt) = LookaheadSet t -> Maybe (LookaheadSet t) -> LookaheadSet t
forall a. a -> Maybe a -> a
fromMaybe LookaheadSet t
forall t. LookaheadSet t
LA.empty (AugNT nt
-> Map (AugNT nt) (LookaheadSet t) -> Maybe (LookaheadSet t)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup AugNT nt
nt Map (AugNT nt) (LookaheadSet t)
fs)
    -- TODO I need a consistent story here of what I define and
    -- export.  FollowSet needs this one below, but you can see the
    -- code duplication with firstsV.  Resolve.

-- | Given a firsts function, find the first set of a list of symbols.
firstsOfVs ::
     Ord t => (AugNT nt -> LookaheadSet t) -> AugVs t nt -> LookaheadSet t
firstsOfVs :: (AugNT nt -> LookaheadSet t) -> AugVs t nt -> LookaheadSet t
firstsOfVs AugNT nt -> LookaheadSet t
firsts AugVs t nt
vs = [LookaheadSet t] -> LookaheadSet t
forall a. Monoid a => [a] -> a
mconcat ([LookaheadSet t] -> LookaheadSet t)
-> [LookaheadSet t] -> LookaheadSet t
forall a b. (a -> b) -> a -> b
$ (V (AugT t) (AugNT nt) -> LookaheadSet t)
-> AugVs t nt -> [LookaheadSet t]
forall a b. (a -> b) -> [a] -> [b]
map V (AugT t) (AugNT nt) -> LookaheadSet t
firstsV' AugVs t nt
vs
  where
    firstsV' :: V (AugT t) (AugNT nt) -> LookaheadSet t
firstsV' (T AugT t
t) = AugT t -> LookaheadSet t
forall t. AugT t -> LookaheadSet t
LA.singleton AugT t
t
    firstsV' (NT AugNT nt
nt) = AugNT nt -> LookaheadSet t
firsts AugNT nt
nt