{-# LANGUAGE TemplateHaskell #-}

module Simple.Plugin(plugin) where

import GHC.Types.Unique.FM
import GHC.Plugins
import qualified GHC.Utils.Error

-- For annotation tests
import Simple.DataStructures

import Control.Monad
import Data.Monoid hiding (Alt)
import Data.Dynamic
import qualified Language.Haskell.TH as TH

plugin :: Plugin
plugin = defaultPlugin {
    installCoreToDos = install,
    pluginRecompile  = purePlugin
  }

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install options todos = do
    putMsgS $ "Simple Plugin Passes Queried"
    putMsgS $ "Got options: " ++ unwords options

    -- Create some actual passes to continue the test.
    return $ CoreDoPluginPass "Main pass" mainPass
             : todos

findNameBinds :: String -> [CoreBind] -> First Name
findNameBinds target = mconcat . map (findNameBind target)

findNameBind :: String -> CoreBind -> First Name
findNameBind target (NonRec b e) = findNameBndr target b
findNameBind target (Rec bes) = mconcat (map (findNameBndr target . fst) bes)

findNameBndr :: String -> CoreBndr -> First Name
findNameBndr target b
  = if getOccString (varName b) == target
    then First (Just (varName b))
    else First Nothing


mainPass :: ModGuts -> CoreM ModGuts
mainPass guts = do
    putMsgS "Simple Plugin Pass Run"
    (_, anns) <- getAnnotations deserializeWithData guts :: CoreM (ModuleEnv [ReplaceWith], NameEnv [ReplaceWith])
    -- Var's have the same uniques as their names. Making a cast from NameEnv to VarEnv safe.
    let anns' = unsafeCastUFMKey anns :: VarEnv [ReplaceWith]
    bindsOnlyPass (mapM (changeBind anns' Nothing)) guts

changeBind :: VarEnv [ReplaceWith] -> Maybe String -> CoreBind -> CoreM CoreBind
changeBind anns mb_replacement (NonRec b e) = changeBindPr anns mb_replacement b e >>= (return . uncurry NonRec)
changeBind anns mb_replacement (Rec bes) = liftM Rec $ mapM (uncurry (changeBindPr anns mb_replacement)) bes

changeBindPr :: VarEnv [ReplaceWith] -> Maybe String -> CoreBndr -> CoreExpr -> CoreM (CoreBndr, CoreExpr)
changeBindPr anns mb_replacement b e = do
    case lookupWithDefaultUFM anns [] b of
        [] -> do
                e' <- changeExpr anns mb_replacement e
                return (b, e')
        [ReplaceWith replace_string] -> do
                e' <- changeExpr anns (Just replace_string) e
                return (b, e')
        _ -> do dflags <- getDynFlags
                error ("Too many change_anns on one binder:" ++ showPpr dflags b)

changeExpr :: VarEnv [ReplaceWith] -> Maybe String -> CoreExpr -> CoreM CoreExpr
changeExpr anns mb_replacement e = let go = changeExpr anns mb_replacement in case e of
        Lit (LitString _) -> case mb_replacement of
                Nothing -> return e
                Just replacement -> do
                        putMsgS "Performing Replacement"
                        return $ Lit (LitString (bytesFS (mkFastString replacement)))
        App e1 e2 -> liftM2 App (go e1) (go e2)
        Lam b e -> liftM (Lam b) (go e)
        Let bind e -> liftM2 Let (changeBind anns mb_replacement bind) (go e)
        Case e b ty alts -> liftM4 Case (go e) (return b) (return ty) (mapM (changeAlt anns mb_replacement) alts)
        Cast e coerce -> liftM2 Cast (go e) (return coerce)
        Tick t e -> liftM (Tick t) (go e)
        _ -> return e

changeAlt :: VarEnv [ReplaceWith] -> Maybe String -> CoreAlt -> CoreM CoreAlt
changeAlt anns mb_replacement (Alt con bs e) = liftM (\e' -> Alt con bs e') (changeExpr anns mb_replacement e)