ghc parser plugin to assign IORefs

VVariable.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}

module VVariable where

import Control.Monad
import Control.Monad.Trans.Writer
import Data.Data.Lens
import Data.Functor.Identity
import Data.Generics
import Data.Monoid
import GHC
import GHC.Parser.PostProcess (mkTokenLocation)
import GHC.Plugins
import GHC.Types.Name.Occurrence

plugin :: Plugin
plugin =
  defaultPlugin
    { parsedResultAction = \_ _ -> (l1 . l2 . l3) (everywhereM (mkM transformHsDo)),
      -- is this traversal in the right direction?
      -- it has to be bottom-up otherwise it will apply the transformation to it's own output,
      -- and I haven't renamed v_x in such a way that this could work.
      pluginRecompile = purePlugin
    }

-- | Lens'
type s :> a = forall f. (Functor f) => (a -> f a) -> s -> f s

l1 :: ParsedResult :> HsParsedModule
l1 f x = (\x' -> x {parsedResultModule = x'}) `fmap` f (parsedResultModule x)

l2 :: HsParsedModule :> Located (HsModule GhcPs)
l2 f x = (\x' -> x {hpm_module = x'}) `fmap` f (hpm_module x)

l3 :: Located e :> e
l3 f (L l e) = L l `fmap` f e

(%~) :: (s :> a) -> (a -> a) -> s -> s
(%~) l f = runIdentity . l (Identity . f)

-- | transformRhs turns
--
-- f v_x v_y
--
-- into
--
-- f !(readIORef v_x) !(readIORef v_y)
--
-- or should it redo MonadicBang? If I review the MonadicBang's
-- source code I should find out how new names are created in
-- the parser. It bangVar :: Has (Uniques :+: Reader DynFlags) sig m => LExpr -> Loc -> m RdrName
transformRhs :: LHsExpr GhcPs -> Hsc (LHsExpr GhcPs)
transformRhs (L l (HsVar _ _ )) = pure _
transformRhs x = pure x

-- | transformHsDo applies transformStmt to every statement in a do block,
-- and the extra do blocks created use the outer do block flavor
transformHsDo :: LHsExpr GhcPs -> Hsc (LHsExpr GhcPs)
transformHsDo (L l (HsDo a b c)) = do
  c' <- mapM (transformStmt (noLocA . HsDo a b . noLocA) . unLoc) (unLoc c)
  return (L l (HsDo a b (noLocA $ map noLocA c')))
transformHsDo x = return x

--  | transformStmt turns
--
--  (v_x, y) <- f
--
--  into
--
--  (_, y) <- do
--        (v_x, y) <- f
--        writeIORef x v_x
--        pure (v_x, y)
transformStmt ::
  ([LStmtLR GhcPs GhcPs (LHsExpr GhcPs)] -> LHsExpr GhcPs) ->
  StmtLR GhcPs idR (LHsExpr GhcPs) ->
  Hsc (StmtLR GhcPs idR (LHsExpr GhcPs))
transformStmt mkdo orig@(BindStmt a lhs rhs) =
  BindStmt a lhs' <$> do
    atVar <- newName "a282393bf"
    pure
      ( mkdo
          ( [noLocA $ BindStmt EpAnnNotUsed (addAt atVar lhs) rhs]
              ++ map mkWriteIORef iorefs
              ++ [returnExpr atVar]
          )
      )
  where
    (lhs', iorefs) = replaceWildP lhs
transformStmt mkdo x = pure x

-- | attAt x y makes x@y
addAt :: OccName -> LPat GhcPs -> LPat GhcPs
addAt n r = noLocA $ AsPat EpAnnNotUsed (noLocA $ mkRdrUnqual n) (L (mkTokenLocation noSrcSpan) HsTok) r

-- | mkWriteIORef (x, y) makes the syntax written as writeIORef x y
--
-- ideally this could be written [| writeIORef $x $y |] or similar
mkWriteIORef :: (RdrName, RdrName) -> LStmtLR GhcPs GhcPs (LHsExpr GhcPs)
mkWriteIORef (x, y) = noLocA $ BodyStmt NoExtField expr noSyntaxExpr noSyntaxExpr
  where
    v = noLocA . HsVar NoExtField . noLocA
    appE a b = noLocA $ HsApp EpAnnNotUsed a b
    infixl 9 `appE`

    expr :: LHsExpr GhcPs
    expr = v writeIORefName `appE` v x `appE` v y

    writeIORefName = mkRdrUnqual (mkVarOcc "writeIORef")

-- it doesn't really have to be a new name it just has to be different
-- than the others bound in the same as-pattern
newName :: String -> Hsc OccName
newName = return . mkVarOcc

returnExpr :: OccName -> LStmtLR GhcPs GhcPs (LHsExpr GhcPs)
returnExpr a = noLocA $ BodyStmt NoExtField (noLocA $ HsVar NoExtField (noLocA $ mkRdrUnqual a)) noSyntaxExpr noSyntaxExpr

-- | extract every variable named v_*, replace it with WildPat, and return the *
replaceWildP :: LPat GhcPs -> (LPat GhcPs, [(RdrName, RdrName)])
replaceWildP = fmap (`appEndo` []) . runWriter . everywhereM (mkM f)
  where
    f :: Pat GhcPs -> Writer (Endo [(RdrName, RdrName)]) (Pat GhcPs)
    f (VarPat a (L x (Unqual b))) | 'v' : '_' : bs <- occNameString b = do
      tell (Endo ((mkRdrUnqual (mkVarOcc bs), mkRdrUnqual b) :))
      pure (WildPat NoExtField)
    f x = pure x