{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Check definition bodies.
--
-- Simple first-order type-checking for the bodies of @define@d constructors.

module BNFC.Check.Expressions ( checkFunction ) where

import BNFC.Prelude
import Lens.Micro.TH (makeLensesFor)

import qualified Data.Map         as Map
import qualified BNFC.Utils.List1 as List1

import qualified BNFC.Abs as A
import BNFC.Abs (HasPosition(..))
import BNFC.CF
import BNFC.Types.Position

import BNFC.Check.Monad


-- | State of the type checker.

data Env = Env
  { _envSignature :: Signature
      -- ^ The function types of labels and definitions.

  , _envContext   :: Context
      -- ^ The base types of definition parameters.

  , _envPosition  :: Position
      -- ^ The current position.
  }
  deriving (Show)

-- | Typing context, mapping parameters to their type.

type Context = Map VarName Type

-- @makeLenses ''Env@ will generate unused definitions if not all lenses are used
makeLensesFor (map (\ s -> ('_':s, s)) ["envSignature", "envContext"]) ''Env

-- | Entry point to check contents of a 'A.Function'.

checkFunction :: Position -> Signature -> LabelName -> [A.Arg] -> A.Exp -> FunType -> Check Function
checkFunction p sig _f params0 exp0 (FunType t ts) = atPosition p $ do
  (params, cxt) <- checkParameters params0 ts `runStateT` mempty
  -- let cxt = Map.fromList $ params <&> \ (Parameter x t) -> (x,t)
  exp    <- checkExp t exp0 `runReaderT` Env sig cxt p
  return $ Function params exp t

-- | Build a typing context from a list of parameters and their types.
--   If the list do not have equal length, report an error and try to
--   patch the situation.

checkParameters :: [A.Arg] -> [Type] -> StateT Context Check [Parameter]
checkParameters = curry $ \case
  ([]  , []) -> return []

  -- More parameters than types: drop the parameters.
  (a:as, []) -> [] <$ do
    atPosition (hasPosition a) $
      recoverableError $ DroppingSpuriousParameters $ a :| as

  -- More types than parameters: make up some dummy parameters.
  ([]  , ts@(_:_)) -> do
    -- If there are not enough parameters given, we can make some up.
    p <- fmap (\ (Position l c) -> (l, c)) <$> askPosition
    let
      xs = [1 .. length ts] <&> \ n -> '_' :| show n
      dummies = map (A.Arg p . A.Identifier . (fromMaybe panicPositionNothing p,) . toList) xs
    recoverableError $ NotEnoughParameters $ List1.fromList xs
    checkParameters dummies ts

  -- Check the first parameter against the first type.
  -- If the parameter is already in the typing context, warn about shadowing.
  -- Add the parameter to the context and continue.

  (A.Arg p (A.Identifier (_, x0)) : as, t:ts) -> atPosition p $ do
    let x = fromMaybe panicEmptyIdentifier $ nonEmpty x0
    -- If x is not lowercase, warn.
    unless (isLower $ List1.head x) $ warn $ ParameterShouldBeLowerCase x
    -- If x is already in the context, warn about shadowing.
    mapM_ (const $ warn $ ShadowingParameter x) =<< gets (Map.lookup x)
    (Parameter x t :) <$> do
      -- Add parameter to context, possibly overwriting the shadowed alter ego.
      modify $ Map.insert x t
      checkParameters as ts

-- | Check a list of expressions against a list of types.
--   The lists are checked for the same length.

checkExps :: LabelName -> [Type] -> [A.Exp] -> ReaderT Env Check [Exp]
checkExps x = curry $ \case
  ([]  , []  ) -> return []
  ([]  , e:es) -> [] <$ do recoverableError $ DroppingSpuriousArguments x $ e :| es
  (t:ts, []  ) -> [] <$ do recoverableError $ MissingArguments x $ t :| ts
  (t:ts, e:es) -> (:) <$> checkExp t e <*> checkExps x ts es

-- | Check an LBNF expression against the given type and return the
--   expression in internal representation.

checkExp :: Type -> A.Exp -> ReaderT Env Check Exp
checkExp t = \case
    A.App p0 (A.Identifier (_, x0)) es0 -> atPosition p0 $ do
      let x = fromMaybe panicEmptyIdentifier $ nonEmpty x0
      warnIfShadowed x
      ft@(FunType t' ts) <- lookupSig x
      inferred t'
      App (labelFromIdentifier x) ft <$> checkExps x ts es0

    A.Var p0 (A.Identifier (_, x0)) -> atPosition p0 $ do
      let x = fromMaybe panicEmptyIdentifier $ nonEmpty x0
      lookupCxt x >>= \case
        Just t' -> inferred t' $> Var (Parameter x t)
        Nothing -> do
          ft@(FunType t' ts) <- lookupSig x
          -- If this is a proper function type, then we are missing arguments
          mapM_ (recoverableError . MissingArguments x) $ nonEmpty ts
          inferred t' $> App (labelFromIdentifier x) ft []

    A.Cons p0 e0 es0 -> atPosition p0 $ do
      elemType $ \ s -> do
        e <- checkExp s e0
        cons s e <$> checkExp t es0

    A.List p0 es0    -> atPosition p0 $ do
      elemType $ \ s -> do
        foldr (cons s) nil <$> mapM (checkExp s) es0

    A.LitInteger _ l   -> inferred tInteger $> LitInteger l
    A.LitDouble  _ l   -> inferred tDouble  $> LitDouble l
    A.LitChar    _ l   -> inferred tChar    $> LitChar l
    A.LitString  _ l   -> inferred tString  $> LitString l
  where
  -- Match inferred type against ascribed type.
  inferred t'
    | t == t'   = return ()
    | otherwise = recoverableError $ ExpectedVsInferredType t t'
  -- Check whether the ascribed type is a list type.
  elemType k =
    case t of
      ListType s -> k s
      BaseType{} -> fatalError $ ListsDontInhabitType t
  -- List constructors
  nil = App LNil (FunType t []) []
  -- Precondition to cons: t == ListT s
  cons s e es = App LCons (FunType t [s,t]) [e,es]

  lookupCxt x = Map.lookup x <$> view envContext
  lookupSig x = do
    (Map.lookup x <$> view envSignature) >>= \case
      Nothing                  -> fatalError $ UndefinedLabel x
      Just (WithPosition _ ft) -> return ft

  -- If a label applied to arguments is one of the parameters,
  -- we issue a warning because it looks confusing, even though
  -- it is not ambiguous, strictly speaking.  Example:
  --
  -- @
  --   define f x = EInt x;
  --   define g f = f f;  -- The first f is shadowed by the parameter f.
  -- @
  warnIfShadowed x =
    mapM_ (const $ warn $ ShadowedByParameter x) =<< lookupCxt x