{-  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 RankNTypes #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Text.GrammarCombinators.Parser.LL1 ( 
  LL1Table(LL1Table),
  calcLL1Table,
  parseLL1
  ) where

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

import Control.Monad
import Control.Monad.State

import Text.GrammarCombinators.Base
import Text.GrammarCombinators.Parser.TopDown

data FirstSet t = 
  FS {
    firstSet :: Set t,
    canBeEmpty :: Bool, 
    canBeEOI :: Bool
    } 

type FirstSetGrammar phi t = forall ix. phi ix -> [FirstSet t]

newtype (Domain phi, Token t) => FSCalculator phi ixT (r :: * -> *) t v = MkFSCalculator {
  calcFS :: FirstSetGrammar phi t -> [FirstSet t]
  }
type FirstSetGrammarRec phi ixT r t rr = forall ix. phi ix -> FSCalculator phi ixT r t (rr ix)

unionL :: (Ord a) => [Set a] -> Set a
unionL = foldr Set.union Set.empty
                                             
instance (Domain phi, Token t) => ProductionRule (FSCalculator phi ixT r t) where
  a >>> b = MkFSCalculator $ \g ->
      do FS fsa ea fa <- calcFS a g
         FS fsb eb fb <- calcFS b g
         return $ FS (if ea then fsa `union` fsb else fsa) (ea && eb) (fa || (ea && fb))
  a ||| b = MkFSCalculator disjFS where
    disjFS :: FirstSetGrammar phi t -> [FirstSet t]
    disjFS g = calcFS a g ++ calcFS b g
  die = MkFSCalculator $ \_ -> [FS Set.empty False False]
  endOfInput = MkFSCalculator $ \_ -> [FS Set.empty False True] 

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

instance (Domain phi, Token t) => LiftableProductionRule (FSCalculator phi ixT r t) where
  epsilonL v _ = epsilon v

instance (Token t, Domain phi) => TokenProductionRule (FSCalculator phi ixT r t) t where
  token c = MkFSCalculator $ \_ -> [FS (singleton c) False False]
  anyToken = MkFSCalculator $ \_ -> [FS (Set.fromList enumerate) False False]
instance (Domain phi, Token t) => RecProductionRule (FSCalculator phi ixT r t) phi r where
  ref idx = MkFSCalculator $ \g -> [FS (unionL $ map firstSet $ g idx) (any canBeEmpty $ g idx) (any canBeEOI $ g idx)]

fixFSGrammar :: (Domain phi, Token t) => FirstSetGrammarRec phi ixT r t rr -> FirstSetGrammar phi t
fixFSGrammar g idx = calcFS (g idx) $ fixFSGrammar g

data (Token t) => LL1Table phi t = LL1Table {
  ruleForTokenTable :: Memo phi (K0 (Map t Int)),
  ruleForEOITable :: Memo phi (K0 (Maybe Int)),
  ruleForEmptyTable :: Memo phi (K0 (Maybe Int))
  }
                                   
calcLL1Table :: forall phi r t rr. (Token t, Domain phi) =>
                GContextFreeGrammar phi t r rr -> LL1Table phi t
calcLL1Table grammar = 
  let
    g :: FirstSetGrammar phi t
    g = fixFSGrammar grammar
    fss :: forall ix. phi ix -> [Set t]
    fss = map firstSet . g
    n :: forall ix. phi ix -> Int
    n = length . fss
    ttableContents :: forall ix. phi ix -> [(t,Int)]
    ttableContents idx = do (fs,i) <- zip (fss idx) [0..n idx-1]
                            c <- Set.toList fs
                            return (c,i)
    rftTable :: forall ix. phi ix -> Map t Int
    rftTable idx = Map.fromListWith notLL1Error $ ttableContents idx
    cbe :: forall ix. phi ix -> [Bool]
    cbe = map canBeEOI . g
    etableContents :: forall ix. phi ix -> [Int]
    etableContents idx = do (True, i) <- zip (cbe idx) [0..n idx-1]
                            return i
    rfeTable :: forall ix. phi ix -> Maybe Int
    rfeTable = listToMaybe . etableContents 
    rfnTable :: forall ix. phi ix -> Maybe Int
    rfnTable = listToMaybe . ntableContents 
    cbn :: forall ix. phi ix -> [Bool]
    cbn = map canBeEmpty . g
    ntableContents :: forall ix. phi ix -> [Int]
    ntableContents idx = do (True, i) <- zip (cbn idx) [0..n idx-1]
                            return i
    notLL1Error = error "Not LL1"
  in LL1Table (toMemoK rftTable) (toMemoK rfeTable) (toMemoK rfnTable)

newtype LLRule phi ixT r t v = MkLLRule {
  llRuleAlts :: [NonBranchingRule phi r t v]
  }

instance Functor (LLRule phi ixT r t) where
  fmap f (MkLLRule rules) = MkLLRule [fmap f rule | rule <- rules]
instance ProductionRule (LLRule phi ixT r t) where
  (MkLLRule rulesa) >>> (MkLLRule rulesb) =
    let
      seqrule = liftM2 ($) 
    in MkLLRule [seqrule rulea ruleb | rulea <- rulesa, ruleb <- rulesb]
  (MkLLRule rulesa) ||| (MkLLRule rulesb) =
    MkLLRule $ rulesa ++ rulesb
  die = MkLLRule []
  endOfInput = MkLLRule [nbrEndOfInput]
  
instance EpsProductionRule (LLRule phi ixT r t) where
  epsilon v = MkLLRule [return v]

instance LiftableProductionRule (LLRule phi ixT r t) where
  epsilonL v _ = MkLLRule [return v]

instance (Token t) => TokenProductionRule (LLRule phi ixT r t) t where
  token t =
    let
      rule = do (c:r) <- MkNBR $ \_ -> get
                if classify c == t 
                  then do MkNBR $ \_ -> put r
                          return c
                  else fail $ errWrongToken c
      errWrongToken c =  show c ++ " read when " ++ show t ++ " expected."
    in MkLLRule [rule]
  anyToken =
    let rule = do (c:r) <- MkNBR $ \_ -> get
                  MkNBR $ \_ -> put r
                  return c
    in MkLLRule [rule]
  
instance RecProductionRule (LLRule phi ixT r t) phi r where
  ref idx = MkLLRule [MkNBR $ \g -> get >>= \s -> unNBR (g idx s) g]

newtype WrapNonBranchingRuleList phi r t ix = WrapNBRL {
  unWrapNBRL :: [NonBranchingRule phi r t (r ix)]
  }

ll1Disambiguate :: forall phi r t. (Domain phi, Token t) =>
                   ProcessingContextFreeGrammar phi t r ->
                   LL1Table phi t ->
                   UnambiguousTopDownGrammar phi r t
ll1Disambiguate gram table =
  let
    tableidx :: phi ix -> K0 (Map t Int) ix
    tableidx = fromMemo (ruleForTokenTable table)
    ttable :: phi ix -> Map t Int
    ttable idx = unK0 $ tableidx idx
    eoitable :: phi ix -> Maybe Int
    eoitable idx = unK0 $ fromMemo (ruleForEOITable table) idx
    emptytable :: phi ix -> Maybe Int
    emptytable idx = unK0 $ fromMemo (ruleForEmptyTable table) idx
    tidx :: phi ix -> ConcreteToken t -> Int
    tidx idx c = fromMaybe (emptyidx idx) $ Map.lookup (classify c) $ ttable idx
    eoiidx :: phi ix -> Int
    eoiidx idx = fromMaybe (emptyidx idx) $ eoitable idx
    emptyidx :: phi ix -> Int
    emptyidx idx = fromJust $ emptytable idx
    candidateRules :: phi ix -> [NonBranchingRule phi r t (r ix)]
    candidateRules idx = llRuleAlts $ gram idx
    memoCR :: phi ix -> [NonBranchingRule phi r t (r ix)]
    memoCR = unWrapNBRL . memoFamily (WrapNBRL . candidateRules)
    ruleForString :: phi ix -> [ConcreteToken t] -> NonBranchingRule phi r t (r ix)
    ruleForString idx (c:_) = memoCR idx !! tidx idx c
    ruleForString idx [] = memoCR idx !! eoiidx idx 
  in ruleForString

parseLL1 :: forall phi ixT r t ix. (Domain phi, Token t, ProductionRule (LLRule phi ixT r t)) =>
            ProcessingContextFreeGrammar phi t r ->
            LL1Table phi t -> phi ix -> [ConcreteToken t] -> Maybe (r ix)
parseLL1 gram table =
  let
    unambGram :: UnambiguousTopDownGrammar phi r t
    unambGram = ll1Disambiguate gram table
  in parseTopDown unambGram