DouKaku?

{-# OPTIONS_GHC -fth #-}

module LocalVariables(localVariables, withLocalVariables) where

import Control.Monad
import Data.Dynamic
import Language.Haskell.TH
import qualified Data.Map as Map

localVariables :: Map.Map String Dynamic
localVariables = error "localVariables outside withLocalVariables"

localVariablesName :: Name
localVariablesName = 'localVariables

withLocalVariables :: Q [Dec] -> Q [Dec]
withLocalVariables = (>>=trDecs)
  where
    trDecs decs = mapM (trDec []) decs

    trDec vars dec = case dec of
      FunD name cls -> liftM (FunD name) $ mapM (trClause vars) cls
      ValD pat body decs -> do
        decs' <- mapM (trDec vars) decs
        body' <- trBody (vars ++ concatMap decVars decs) body
        return $ ValD pat body' decs'
      _ -> return dec

    trClause vars (Clause pats body decs) = do
      decs' <- mapM (trDec vars) decs
      body' <- trBody (vars ++ concatMap patVars pats ++ concatMap decVars decs) body
      return $ Clause pats body' decs'

    trBody vars (GuardedB ges) = liftM GuardedB $ mapM trGuarded ges
      where
        trGuarded (guard, exp) = do
          exp' <- trExp vars exp
          return (guard, exp')
    trBody vars (NormalB exp) = liftM NormalB $ trExp vars exp

    trExp :: [Name] -> Exp -> Q Exp
    trExp vars exp = case exp of
      VarE name
        | name == localVariablesName -> replacement vars
        | otherwise -> return exp
      AppE x y -> liftM2 AppE (rec x) (rec y)
      InfixE x op y -> liftM3 InfixE (maybeMapM rec x) (rec op) (maybeMapM rec y)
      LamE pats e -> liftM (LamE pats) $ trExp (vars ++ concatMap patVars pats) e
      TupE es -> liftM TupE $ mapM rec es
      CondE c t f -> liftM3 CondE (rec c) (rec t) (rec f)
      LetE decs e -> do
        decs' <- mapM (trDec vars) decs
        e' <- trExp (vars ++ concatMap decVars decs) e
        return $ LetE decs' e'
      CaseE e matches -> liftM2 CaseE (rec e) (mapM (trMatch vars) matches)
      DoE ss -> liftM DoE $ trStmts vars ss
      CompE ss -> liftM CompE $ trStmts vars ss
      ArithSeqE rng -> liftM ArithSeqE $ trRange vars rng
      ListE es -> liftM ListE $ mapM rec es
      SigE e tp -> liftM2 SigE (rec e) (return tp)
      RecConE name fexps -> liftM (RecConE name) $ mapM (trFexp vars) fexps
      RecUpdE e fexps -> liftM2 RecUpdE (rec e) (mapM (trFexp vars) fexps)
      _ -> return exp
      where
        rec e = trExp vars e
        maybeMapM f Nothing = return Nothing
        maybeMapM f (Just v) = liftM Just $ f v

    trFexp vars (name, e) = do
      e' <- trExp vars e
      return (name, e')

    trRange vars rng = case rng of
      FromR x -> liftM FromR $ tr x
      FromThenR x y -> liftM2 FromThenR (tr x) (tr y)
      FromToR x y -> liftM2 FromToR (tr x) (tr y)
      FromThenToR x y z -> liftM3 FromThenToR (tr x) (tr y) (tr z)
      where
        tr e = trExp vars e

    trStmts :: [Name] -> [Stmt] -> Q [Stmt]
    trStmts vars [] = return []
    trStmts vars (stmt:rest) = case stmt of
      BindS pat exp -> let vars' = vars ++ patVars pat
        in liftM2 (:) (liftM (BindS pat) (trExp vars' exp)) (trStmts vars' rest)
      LetS decs -> liftM2 (:) (liftM LetS $ mapM (trDec vars) decs) (trStmts (vars ++ concatMap decVars decs) rest)
      NoBindS exp -> liftM2 (:) (liftM NoBindS $ trExp vars exp) (trStmts vars rest)
      ParS _ -> error "ParS: what's this?"

    trMatch vars (Match pat body decs) = do
      decs' <- mapM (trDec vars) decs
      body' <- trBody (vars ++ patVars pat ++ concatMap decVars decs) body
      return $ Match pat body' decs'

    decVars dec = case dec of
      FunD name _ -> [name]
      ValD pat _ _ -> patVars pat
      _ -> []

    patVars pat = case pat of
      VarP name -> [name]
      TupP pats -> concatMap patVars pats
      ConP name pats -> concatMap patVars pats
      InfixP p0 _ p1 -> patVars p0 ++ patVars p1
      TildeP p -> patVars p
      AsP name p -> name : patVars p
      RecP _ fps -> concatMap fpatVars fps
      ListP pats -> concatMap patVars pats
      _ -> []

    fpatVars (_, pat) = patVars pat

    replacement vars = [| Map.fromList $(list) |]
      where
        list = listE $ map entry vars
        entry name = [| ($(str), toDyn $(obj)) |]
          where
            str = return $ LitE $ StringL $ show name
            obj = return $ VarE name

{-
  使い方：
    別モジュールで

    import LocalVariables

    $(withLocalVariables [d|

      foo :: IO (Map.Map String Dynamic)
      foo = do
        let x = 1 :: Int
        let y = "hello"
        return localVariables

      |])

    のようにすると、localVariablesの出現を置換した上でfooが定義される
-}