{-  Copyright 2010 Dominique Devriese

    This file is part of the grammar-combinators library.

    The grammar-combinators library is free software: you can
    redistribute it and/or modify it under the terms of the GNU
    Lesser General Public License as published by the Free
    Software Foundation, either version 3 of the License, or (at
    your option) any later version.

    Foobar is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General
    Public License along with Foobar. If not, see
    <http://www.gnu.org/licenses/>.
-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}

module Text.GrammarCombinators.Utils.CalcFirst (
  FirstSet (FS, firstSet, canBeEmpty, canBeEOI),
  FSCalculator, FirstSetGrammar, calcFS,
  calcFirst
  ) where

import Data.Set (Set, union, singleton)
import Data.Enumerable (enumerate)
import qualified Data.Set as Set

import Text.GrammarCombinators.Base

data FirstSet t = 
  FS {
    firstSet :: Set t,
    canBeEmpty :: Bool, 
    canBeEOI :: Bool
    } deriving (Show)

newtype FSCalculator phi (r :: * -> *) t rr v = MkFSCalculator {
  calcFS :: FirstSetGrammar phi r t rr -> FirstSet t
  }
type FirstSetGrammar phi r t rr =
  forall ix. phi ix -> FSCalculator phi r t rr (rr ix)

instance (Token t) => ProductionRule (FSCalculator phi r t rr) where
  a >>> b = MkFSCalculator $ \g ->
    let FS fsa ea fa = calcFS a g
        FS fsb eb fb = calcFS b g
    in FS (if ea then fsa `union` fsb else fsa) (ea && eb) (fa || (ea && fb))
  a ||| b = MkFSCalculator $ \g ->
    let fa = calcFS a g
        fb = calcFS b g
    in FS { 
      firstSet = firstSet fa `union` firstSet fb, 
      canBeEmpty = canBeEmpty fa || canBeEmpty fb,
      canBeEOI = canBeEOI fa || canBeEOI fb }
  die = MkFSCalculator $ \_ -> FS Set.empty False False
  endOfInput = MkFSCalculator $ \_ -> FS Set.empty False True 

instance (Token t) => EpsProductionRule (FSCalculator phi r t rr) where
  epsilon _ = MkFSCalculator $ \_ -> FS Set.empty True False

instance (Token t) => LiftableProductionRule (FSCalculator phi r t rr) where
  epsilonL _ _ = MkFSCalculator $ \_ -> FS Set.empty True False

newtype WrapFSC phi r t rr ix = WFSC { unWFSC :: FSCalculator phi r t rr (rr ix) }

blockRecurse :: (EqFam phi, Token t) =>
                FirstSetGrammar phi r t rr -> phi idx -> FirstSetGrammar phi r t rr
blockRecurse gram idx = unWFSC . overrideIdx (WFSC . gram) idx (WFSC die)

instance (Token t) => TokenProductionRule (FSCalculator phi r t rr) t where
  token c = MkFSCalculator $ \_ -> FS (singleton c) False False
  anyToken = MkFSCalculator $ \_ -> FS allTokens False False
    where allTokens = Set.fromList enumerate 
      
instance (Token t, EqFam phi) => RecProductionRule (FSCalculator phi r t rr) phi r where
  ref idx = MkFSCalculator $ \g -> calcFS (g idx) $ blockRecurse g idx
instance (Token t, EqFam phi) => LoopProductionRule (FSCalculator phi r t rr) phi r where
  -- first-set-wise, many1ref is identical to ref
  -- manyRef will be defaulted correctly
  many1Ref idx = MkFSCalculator $ \g -> calcFS (ref idx) g 

calcFirst :: (Domain phi, Token t) =>
                GExtendedContextFreeGrammar phi t r rr ->
                phi ix -> FirstSet t
calcFirst g idx = calcFS (ref idx) g