content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE RecordWildCards #-}
  
  -- | Type checker for defined syntax constructors @define f xs = e@.
  
  module BNFC.TypeChecker
    ( -- * Type checker entry point
      runTypeChecker
    , checkDefinitions
      -- * Backdoor for rechecking defined syntax constructors for list types
    , checkDefinition'
    , buildSignature, buildContext, ctxTokens, isToken
    , ListConstructors(..)
    ) where
  
  import Control.Monad
  import Control.Monad.Except (MonadError(..))
  import Control.Monad.Reader
  
  import Data.Bifunctor
  import Data.Char
  import Data.Either (partitionEithers)
  
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  
  import BNFC.CF
  import BNFC.PrettyPrint
  
  -- * Error monad
  
  type TCError = WithPosition String
  
  -- | Type checking monad, reports errors.
  newtype Err a = Err { unErr :: ReaderT Position (Either TCError) a }
    deriving (Functor, Applicative, Monad, MonadReader Position)
  
  instance MonadError String Err where
    throwError msg = Err $ do
      pos <- ask
      throwError $ WithPosition pos msg
    catchError m h = Err $ do
      unErr m `catchError` \ (WithPosition _ msg) -> unErr (h msg)
  
  withPosition :: Position -> Err a -> Err a
  withPosition pos = local (const pos)
  
  runTypeChecker :: Err a -> Either String a
  runTypeChecker m = first blendInPosition $ unErr m `runReaderT` NoPosition
  
  -- * Types and context
  
  data Context = Ctx
    { ctxLabels :: Signature         -- ^ Types of labels, extracted from rules.
    , ctxTokens :: [String]          -- ^ User-defined token types.
    , ctxLocals :: Telescope         -- ^ Types of local variables of a definition.
    }
  
  data ListConstructors = LC
    { nil   :: Base -> (String, Type)  -- ^ 'Base' is the element type. 'Type' the list type.
    , cons  :: Base -> (String, Type)
    }
  
  dummyConstructors :: ListConstructors
  dummyConstructors = LC
    { nil  = \ b -> ("[]" , FunT [] (ListT b))
    , cons = \ b -> ("(:)", FunT [b, ListT b] (ListT b))
    }
  
  -- * Type checker for definitions and expressions
  
  -- | Entry point.
  checkDefinitions :: CF -> Err CF
  checkDefinitions cf = do
    let ctx = buildContext cf
    let (pragmas, defs0) = partitionEithers $ map isFunDef $ cfgPragmas cf
    defs <- mapM (checkDefinition ctx) defs0
    return cf { cfgPragmas = pragmas ++ map FunDef defs }
  
  checkDefinition :: Context -> Define -> Err Define
  checkDefinition ctx (Define f args e0 _) = do
    let xs = map fst args  -- Throw away dummy types.
    (tel, (e, b)) <- checkDefinition' dummyConstructors ctx f xs e0
    return $ Define f tel e b
  
  checkDefinition'
    :: ListConstructors  -- ^ Translation of the list constructors.
    -> Context           -- ^ Signature (types of labels).
    -> RFun              -- ^ Function name.
    -> [String]          -- ^ Function arguments.
    -> Exp               -- ^ Function body.
    -> Err (Telescope, (Exp, Base))  -- ^ Typed arguments, translated body, type of body.
  checkDefinition' list ctx ident xs e =
    withPosition (wpPosition ident) $
      do  unless (isLower $ head f) $ throwError $
            "Defined functions must start with a lowercase letter."
          t@(FunT ts t') <- lookupCtx f ctx `catchError` \_ ->
                                  throwError $ "'" ++ f ++ "' must be used in a rule."
          let expect = length ts
              given  = length xs
          unless (expect == given) $ throwError $ concat
            [ "'", f, "' is used with type ", show t
            , " but defined with ", show given, " argument", plural given ++ "."
            ]
          e' <- checkExp list (setLocals ctx $ zip xs ts) e t'
          return (zip xs ts, (e', t'))
      `catchError` \ err -> throwError $
        "In the definition " ++ unwords (f : xs ++ ["=", prettyShow e, ";"]) ++ "\n  " ++ err
      where
          f = wpThing ident
          plural 1 = ""
          plural _ = "s"
  
  checkExp :: ListConstructors -> Context -> Exp -> Base -> Err Exp
  checkExp list ctx = curry $ \case
    (App "[]" _ []     , ListT t      ) -> return (uncurry App (nil list t) [])
    (App "[]" _ _      , _            ) -> throwError $
      "[] is applied to too many arguments."
  
    (App "(:)" _ [e,es], ListT t      ) -> do
      e'  <- checkExp list ctx e t
      es' <- checkExp list ctx es (ListT t)
      return $ uncurry App (cons list t) [e',es']
  
    (App "(:)" _ es  , _              ) -> throwError $
      "(:) takes 2 arguments, but has been given " ++ show (length es) ++ "."
  
    (e@(App x _ es)  , t              ) -> checkApp e x es t
    (e@(Var x)       , t              ) -> e <$ checkApp e x [] t
    (e@LitInt{}      , BaseT "Integer") -> return e
    (e@LitDouble{}   , BaseT "Double" ) -> return e
    (e@LitChar{}     , BaseT "Char"   ) -> return e
    (e@LitString{}   , BaseT "String" ) -> return e
    (e               , t              ) -> throwError $
      prettyShow e ++ " does not have type " ++ show t ++ "."
    where
    checkApp e x es t = do
      ft@(FunT ts t') <- lookupCtx x ctx
      es' <- matchArgs ts
      unless (t == t') $ throwError $ prettyShow e ++ " has type " ++ show t' ++ ", but something of type " ++ show t ++ " was expected."
      return $ App x ft es'
      where
      matchArgs ts
        | expect /= given   = throwError $ "'" ++ x ++ "' takes " ++ show expect ++ " arguments, but has been given " ++ show given ++ "."
        | otherwise         = zipWithM (checkExp list ctx) es ts
        where
          expect = length ts
          given  = length es
  
  -- * Context handling
  
  -- | Create context containing the types of all labels,
  --   computed from the rules.
  --
  --   Fail if a label is used at different types.
  --
  buildSignature :: [Rule] -> Err Signature
  buildSignature rules = do
    -- Build label signature with duplicates
    let sig0 = Map.fromListWith mappend $ map (second Set.singleton) labels
    -- Check for duplicates; extract from singleton sets.
    sig <- forM (Map.toAscList sig0) $ \ (f,ts) ->
      case Set.toList ts of
        []  -> undefined  -- impossible
        [t] -> return (f,t)
        ts' -> throwError $ unlines $ concat
          [ [ "The label '" ++ f ++ "' is used at conflicting types:" ]
          , map (("  " ++) . blendInPosition . fmap show) ts'
          ]
    return $ Map.fromAscList sig
    where
      mkType cat args = FunT [ mkBase t | Left t <- args ]
                             (mkBase cat)
      mkBase t
          | isList t  = ListT $ mkBase $ normCatOfList t
          | otherwise = BaseT $ catToStr $ normCat t
  
      labels =
        [ (x, WithPosition pos $ mkType (wpThing cat) args)
          | Rule f@(WithPosition pos x) cat args _ <- rules
          , not (isCoercion f)
          , not (isNilCons f)
        ]
  
  buildContext :: CF -> Context
  buildContext cf = Ctx
      { ctxLabels = cfgSignature cf
      , ctxTokens = ("Ident" : tokenNames cf)
      , ctxLocals = []
      }
  
  isToken :: String -> Context -> Bool
  isToken x ctx = elem x $ ctxTokens ctx
  
  setLocals :: Context -> [(String,Base)] -> Context
  setLocals ctx xs = ctx { ctxLocals = xs }
  
  lookupCtx :: String -> Context -> Err Type
  lookupCtx x ctx
    | isToken x ctx = return $ FunT [BaseT "String"] (BaseT x)
    | otherwise     = do
      case lookup x $ ctxLocals ctx of
        Just b -> return $ FunT [] b
        Nothing -> do
          case Map.lookup x $ ctxLabels ctx of
            Nothing -> throwError $ "Undefined symbol '" ++ x ++ "'."
            Just t  -> return $ wpThing t
  