content: {-# LANGUAGE CPP               #-}
  {-# LANGUAGE OverloadedStrings #-}
  -- | Module for macro expansion.
  module Language.Finkel.Expand
    ( expand
    , expand1
    , expands
    , expands'
    ) where
  
  #include "ghc_modules.h"
  
  -- base
  import           Data.Char                       (isLower)
  import           Data.Foldable                   (foldlM)
  
  -- containers
  import qualified Data.Map                        as Map
  
  -- ghc
  import           GHC_Types_SrcLoc                (GenLocated (..))
  import           GHC_Utils_Outputable            (Outputable (..), SDoc, cat,
                                                    fsep, nest, vcat)
  
  -- Internal
  import           Language.Finkel.Data.FastString (FastString, unconsFS)
  import           Language.Finkel.Fnk
  import           Language.Finkel.Form
  
  
  -- ---------------------------------------------------------------------
  --
  -- Macro expander
  --
  -- ---------------------------------------------------------------------
  
  -- | Returns a list of bounded names in let expression.
  boundedNames :: Code -> [FastString]
  boundedNames form =
    case unCode form of
      List xs          -> concatMap boundedName xs
      Atom (ASymbol n) -> [n]
      _                -> []
  {-# INLINABLE boundedNames #-}
  
  boundedName :: Code -> [FastString]
  boundedName form =
    case unCode form of
      List (LForm (L _ (Atom (ASymbol "="))):n:_) -> boundedNameOne n
      _                                           -> []
  {-# INLINABLE boundedName #-}
  
  boundedNameOne :: Code -> [FastString]
  boundedNameOne form =
    case unCode form of
      Atom (ASymbol n) -> [n]
      List ns          -> concatMap f ns
      HsList ns        -> concatMap f ns
      _                -> []
    where
      f x =
        case unCode x of
          Atom (ASymbol n) | startsWithLower n -> [n]
          _                                    -> []
  {-# INLINABLE boundedNameOne #-}
  
  startsWithLower :: FastString -> Bool
  startsWithLower fs = case unconsFS fs of
    Just (c, _) -> isLower c
    _           -> False
  {-# INLINABLE startsWithLower #-}
  
  -- | Perform 'Fnk' action with temporary shadowed macro environment.
  withShadowing :: [FastString] -- ^ Names of macro to shadow.
                -> Fnk a -- ^ Action to perform.
                -> Fnk a
  withShadowing toShadow act = do
    fnk_env <- getFnkEnv
    let emacros = envMacros fnk_env
        tmacros = envTmpMacros fnk_env
        f name _ = unMacroName name `notElem` toShadow
    putFnkEnv fnk_env { envMacros = Map.filterWithKey f emacros
                      , envTmpMacros = map (Map.filterWithKey f) tmacros }
    result <- act
    modifyFnkEnv (\e -> e { envMacros = emacros
                          , envTmpMacros = tmacros })
    return result
  
  -- | Expand forms, with taking care of @begin@ special form.
  expands :: [Code] -> Fnk [Code]
  expands forms = do
    fnk_env <- getFnkEnv
    let macro_names me =
          if null me
             then nest 2 "None"
             else nest 2 (fsep (map (ppr . unMacroName) (Map.keys me)))
        tmp_macros = Map.unions (envTmpMacros fnk_env)
    debug fnk_env
          Nothing
          [ "Global macros:",  macro_names (envMacros fnk_env)
          , "Temporary macros:", macro_names tmp_macros ]
    expands' forms
  
  -- | Internal works for 'expands'.
  expands' :: [Code] -> Fnk [Code]
  expands' = fmap concat . mapM expand'
  {-# INLINABLE expands' #-}
  
  -- | Expand form to list of 'Code', supports special form /begin/.
  expand' :: Code -> Fnk [Code]
  expand' form =
    case unCode form of
      List (hd:_) | Atom (ASymbol ":quote") <- unCode hd -> return [form]
      _ -> do
       form' <- expand form
       case unCode form' of
         List (LForm (L _ (Atom (ASymbol ":begin"))):rest) ->
           case rest of
             [] -> return []
             _  -> expands' rest
         _ -> return [form']
  {-# INLINABLE expand' #-}
  
  -- | Recursively expands the given 'Code'.
  expand :: Code -> Fnk Code
  expand form =
    case unLForm form of
      L l (List forms) ->
        case forms of
          -- Expand `let' expression, `do' expression, `case' expression, lambda
          -- expression and function binding with shadowing the lexically bounded
          -- names. Expansion of other forms are done without name shadowing.
          -- This function does not expand quoted forms, to preserve the structure
          -- of the quoted forms containing `:begin'.
          kw@(LForm (L _ (Atom (ASymbol x)))):y:rest
            | x == "let"            -> expandLet l kw y rest
            | x == "do"             -> expandDo l kw (y:rest)
            | x == "case"           -> expandCase l kw y rest
            | x == "where"          -> expandWhere l kw y rest
            | x == "=" || x == "\\" -> expandFunBind l kw (y:rest)
            | x == ":quote"         -> return form
          _                         -> expandList l forms
  
      L l (HsList forms) ->
        -- Without recursively calling 'expand' on the result, cannot expand
        -- macro-generating macros.
        LForm . L l . HsList <$> expands' forms
  
      -- Non-list forms are untouched.
      _ -> return form
    where
      expandLet l kw binds body = do
        binds' <- expand binds
        let bounded = boundedNames binds'
        body' <- withShadowing bounded (expands' body)
        return $! LForm (L l (List (kw:binds':body')))
  
      expandDo l kw body = do
        (_, body') <- foldlM expandInDo ([], []) body
        return $! LForm (L l (List (kw:reverse body')))
  
      expandFunBind l kw rest = do
        let args = init rest
            body = last rest
            bounded = concatMap boundedNameOne args
        args' <- expands' args
        body' <- withShadowing bounded (expand body)
        return $! LForm (L l (List (kw:args'++[body'])))
  
      expandCase l kw expr rest = do
        let go acc xs =
              case xs of
                -- Pattern may have prefix (e.g. '~' for lazy pattern)
                pat_prefix:pat:expr0:rest0
                  | pat_prefix == tilde -> do
                    pat' <- expand pat
                    expr1 <- withShadowing (boundedNameOne pat')
                                           (expand expr0)
                    go (expr1:pat':pat_prefix:acc) rest0
                pat:expr0:rest0 -> do
                  pat' <- expand pat
                  expr1 <- withShadowing (boundedNameOne pat')
                                         (expand expr0)
                  go (expr1:pat':acc) rest0
                _               -> return acc
            tilde = LForm (L l (Atom (ASymbol "~")))
        expr' <- expand expr
        rest' <- go [] rest
        return $! LForm (L l (List (kw:expr':reverse rest')))
  
      expandWhere l kw expr rest = do
        rest' <- expands' rest
        let bounded = concatMap boundedName rest'
        expr' <- withShadowing bounded (expand expr)
        return $! LForm (L l (List (kw:expr':rest')))
  
      expandList l forms =
        case forms of
          sym@(LForm (L _ (Atom (ASymbol k)))) : rest -> do
            fnk_env <- getFnkEnv
            case lookupMacro k fnk_env of
              Just m  -> do_expand k (macroFunction m) >>= expand
              Nothing -> LForm . L l . List . (sym:) <$> expands' rest
          _ -> do
            forms' <- expands' forms
            return $! LForm (L l (List forms'))
  
      do_expand k f =
        do fnk_env <- getFnkEnv
           debug fnk_env Nothing [vcat ["Expanding:", nest 2 (ppr form)]]
           ret0 <- f form
           debug fnk_env Nothing [cat [ppr k, " ==>"], nest 2 (ppr ret0)]
           return ret0
  
  expandInDo ::
     ([FastString], [Code]) -> Code -> Fnk ([FastString], [Code])
  expandInDo (bounded, xs) x = do
    let newbind =
          case x of
            LForm (L _ (List (LForm (L _ (Atom (ASymbol sym))):n:_)))
              | sym == "<-" -> boundedNameOne n
            _               -> []
    x' <- withShadowing bounded (expand x)
    return (newbind ++ bounded, x':xs)
  {-# INLINABLE expandInDo #-}
  
  -- | Expand given form once if the form is a macro form, otherwise
  -- return the given form.
  expand1 :: Code -> Fnk Code
  expand1 form =
    case unLForm form of
      L _l (List ((LForm (L _ (Atom (ASymbol k)))) : _)) -> do
        fnk_env <- getFnkEnv
        case lookupMacro k fnk_env of
          Just m  -> macroFunction m form
          Nothing -> return form
      _ -> return form
  {-# INLINABLE expand1 #-}
  
  -- | Debug function fot macro expansion.
  debug :: FnkEnv -> Maybe SDoc -> [SDoc] -> Fnk ()
  debug = debugWith Fnk_trace_expand
  {-# INLINABLE debug #-}
  
  debugWith :: FnkDebugFlag -> FnkEnv -> Maybe SDoc -> [SDoc] -> Fnk ()
  debugWith debug_flag fnk_env mb_extra msgs0 =
    let msgs1 = maybe msgs0 (: msgs0) mb_extra
    in  debugWhen fnk_env debug_flag msgs1
  {-# INLINABLE debugWith #-}
  