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

module Data.Cfg.Internal.FollowSet
  ( followSetMap
  ) where

import Control.Monad (guard)
import Data.Cfg.Augment
import Data.Cfg.Cfg
import Data.Cfg.Collect (collectOnFirst)
import Data.Cfg.FixedPoint (fixedPoint)
import Data.Cfg.Internal.FirstSet (firstsOfVs)
import Data.Cfg.LookaheadSet hiding (unions)
import qualified Data.Cfg.LookaheadSet as LA
import Data.List (tails)
import qualified Data.Map.Strict as M
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import qualified Data.Set as S

-- | Represents the environment following a nonterminal symbol.  A
-- production @foo ::= <vs> bar <vs'>@ will contribute a 'FollowSite' record
-- with @ntTail == <vs'>@ and @prodHead == foo@, where @<vs>@ is a
-- (possibly empty) list of vocabulary symbols.
data FollowSite t nt = FollowSite
  { FollowSite t nt -> AugVs t nt
ntTail :: AugVs t nt
  , FollowSite t nt -> AugNT nt
prodHead :: AugNT nt
  }

-- | Calculates a map that gives all the follow sites in the grammar
-- for the given nonterminal.
followSitesMap ::
     (Cfg cfg (AugT t) (AugNT nt), Ord nt)
  => cfg (AugT t) (AugNT nt)
  -> M.Map (AugNT nt) [FollowSite t nt]
followSitesMap :: cfg (AugT t) (AugNT nt) -> Map (AugNT nt) [FollowSite t nt]
followSitesMap cfg (AugT t) (AugNT nt)
cfg =
  [(AugNT nt, [FollowSite t nt])] -> Map (AugNT nt) [FollowSite t nt]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(AugNT nt, [FollowSite t nt])]
 -> Map (AugNT nt) [FollowSite t nt])
-> ([(AugNT nt, FollowSite t nt)]
    -> [(AugNT nt, [FollowSite t nt])])
-> [(AugNT nt, FollowSite t nt)]
-> Map (AugNT nt) [FollowSite t nt]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(AugNT nt, FollowSite t nt)] -> [(AugNT nt, [FollowSite t nt])]
forall a b. Eq a => [(a, b)] -> [(a, [b])]
collectOnFirst ([(AugNT nt, FollowSite t nt)] -> Map (AugNT nt) [FollowSite t nt])
-> [(AugNT nt, FollowSite t nt)]
-> Map (AugNT nt) [FollowSite t nt]
forall a b. (a -> b) -> a -> b
$ do
    AugNT nt
prodHd <- 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
prodHd
    Bool -> [()]
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (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)
    Vs (AugT t) (AugNT nt)
rhs <- [Vs (AugT t) (AugNT nt)]
rhss
    NT AugNT nt
nt:Vs (AugT t) (AugNT nt)
tl <- Vs (AugT t) (AugNT nt) -> [Vs (AugT t) (AugNT nt)]
forall a. [a] -> [[a]]
tails Vs (AugT t) (AugNT nt)
rhs
    (AugNT nt, FollowSite t nt) -> [(AugNT nt, FollowSite t nt)]
forall (m :: * -> *) a. Monad m => a -> m a
return (AugNT nt
nt, FollowSite :: forall t nt. AugVs t nt -> AugNT nt -> FollowSite t nt
FollowSite {ntTail :: Vs (AugT t) (AugNT nt)
ntTail = Vs (AugT t) (AugNT nt)
tl, prodHead :: AugNT nt
prodHead = AugNT nt
prodHd})

-- | Given what we know of firsts and follows, find the first set of a
-- follow site.
firstsOfFollowSite ::
     forall t nt. (Ord t, Ord nt)
  => (AugNT nt -> LookaheadSet t)
  -> M.Map (AugNT nt) (LookaheadSet t)
  -> FollowSite t nt
  -> LookaheadSet t
firstsOfFollowSite :: (AugNT nt -> LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
-> FollowSite t nt
-> LookaheadSet t
firstsOfFollowSite AugNT nt -> LookaheadSet t
firsts Map (AugNT nt) (LookaheadSet t)
knownFollows FollowSite t nt
followSite =
  LookaheadSet t
firstsOfNTTail LookaheadSet t -> LookaheadSet t -> LookaheadSet t
forall a. Semigroup a => a -> a -> a
<> LookaheadSet t
firstsOfProdHead
  where
    firstsOfNTTail, firstsOfProdHead :: LookaheadSet t
    firstsOfNTTail :: LookaheadSet t
firstsOfNTTail = (AugNT nt -> LookaheadSet t) -> AugVs t nt -> LookaheadSet t
forall t nt.
Ord t =>
(AugNT nt -> LookaheadSet t) -> AugVs t nt -> LookaheadSet t
firstsOfVs AugNT nt -> LookaheadSet t
firsts (FollowSite t nt -> AugVs t nt
forall t nt. FollowSite t nt -> AugVs t nt
ntTail FollowSite t nt
followSite)
    firstsOfProdHead :: LookaheadSet t
firstsOfProdHead = Map (AugNT nt) (LookaheadSet t)
knownFollows Map (AugNT nt) (LookaheadSet t) -> AugNT nt -> LookaheadSet t
forall k a. Ord k => Map k a -> k -> a
M.! FollowSite t nt -> AugNT nt
forall t nt. FollowSite t nt -> AugNT nt
prodHead FollowSite t nt
followSite

-- | Returns the follow sets for the grammar as a map.
followSetMap ::
     forall cfg t nt. (Cfg cfg (AugT t) (AugNT nt), Ord nt, Ord t)
  => cfg (AugT t) (AugNT nt)
             -- ^ the grammar
  -> (AugNT nt -> LookaheadSet t)
             -- ^ 'firstSet' for the grammar
  -> M.Map (AugNT nt) (LookaheadSet t)
followSetMap :: cfg (AugT t) (AugNT nt)
-> (AugNT nt -> LookaheadSet t) -> Map (AugNT nt) (LookaheadSet t)
followSetMap cfg (AugT t) (AugNT nt)
cfg AugNT nt -> LookaheadSet t
fs = (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)
initMap
  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)
oldFols = (AugNT nt -> LookaheadSet t -> LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey (\AugNT nt
k LookaheadSet t
v -> [LookaheadSet t] -> LookaheadSet t
forall t. Ord t => [LookaheadSet t] -> LookaheadSet t
LA.unions ([LookaheadSet t] -> LookaheadSet t)
-> [LookaheadSet t] -> LookaheadSet t
forall a b. (a -> b) -> a -> b
$ AugNT nt -> LookaheadSet t -> [LookaheadSet t]
f AugNT nt
k LookaheadSet t
v) Map (AugNT nt) (LookaheadSet t)
oldFols
      where
        f :: AugNT nt -> LookaheadSet t -> [LookaheadSet t]
        f :: AugNT nt -> LookaheadSet t -> [LookaheadSet t]
f AugNT nt
nt LookaheadSet t
oldFollows =
          LookaheadSet t
oldFollows LookaheadSet t -> [LookaheadSet t] -> [LookaheadSet t]
forall a. a -> [a] -> [a]
: (FollowSite t nt -> LookaheadSet t)
-> [FollowSite t nt] -> [LookaheadSet t]
forall a b. (a -> b) -> [a] -> [b]
map ((AugNT nt -> LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
-> FollowSite t nt
-> LookaheadSet t
forall t nt.
(Ord t, Ord nt) =>
(AugNT nt -> LookaheadSet t)
-> Map (AugNT nt) (LookaheadSet t)
-> FollowSite t nt
-> LookaheadSet t
firstsOfFollowSite AugNT nt -> LookaheadSet t
fs Map (AugNT nt) (LookaheadSet t)
oldFols) [FollowSite t nt]
folSites
          where
            folSites :: [FollowSite t nt]
folSites = [FollowSite t nt]
-> AugNT nt
-> Map (AugNT nt) [FollowSite t nt]
-> [FollowSite t nt]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [] AugNT nt
nt Map (AugNT nt) [FollowSite t nt]
followSitesMap'
    initMap :: M.Map (AugNT nt) (LookaheadSet t)
    initMap :: Map (AugNT nt) (LookaheadSet t)
initMap =
      [(AugNT nt, LookaheadSet t)] -> Map (AugNT nt) (LookaheadSet t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        [ ( AugNT nt
nt
          , case AugNT nt
nt of
              AugNT nt
StartSymbol -> AugT t -> LookaheadSet t
forall t. AugT t -> LookaheadSet t
singleton AugT t
forall t. AugT t
EOF
              AugNT nt
_ -> LookaheadSet t
forall t. LookaheadSet t
empty)
        | AugNT nt
nt <- [AugNT nt]
nts
        ]
      where
        nts :: [AugNT nt]
nts = 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
    followSitesMap' :: M.Map (AugNT nt) [FollowSite t nt]
    followSitesMap' :: Map (AugNT nt) [FollowSite t nt]
followSitesMap' = cfg (AugT t) (AugNT nt) -> Map (AugNT nt) [FollowSite t nt]
forall (cfg :: * -> * -> *) t nt.
(Cfg cfg (AugT t) (AugNT nt), Ord nt) =>
cfg (AugT t) (AugNT nt) -> Map (AugNT nt) [FollowSite t nt]
followSitesMap cfg (AugT t) (AugNT nt)
cfg