Haskell/XMonad: wrapper around a Monad that also keeps track of data

Question

This is a followup to Ben's previous answer. I had asked for type checking for cases in which X t actions "require cleanup" (ungrabbing of buttons and/or keyboard after it has been completed). His response was a monadic wrapper NeedsCleanup, for which my current implementation goes something like this:

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

newtype NeedsCleanup m t = 
  NeedsCleanup
    {
      -- | Escape hatch from the NeedsCleanup Monad;
      --   returns the original action.
      original_action :: m t
    }
  deriving (Functor, Applicative, Monad)

-- | executes unclean_action and cleans up afterwards.
--   (cleanedUp action) is a normal X() action
cleanedUp :: NeedsCleanup X t -> X t
cleanedUp unclean_action = do
  result <- original_action unclean_action
  doCleanup
  return result

That way, if action has type NeedsCleanup X (), I can't accidentally use it as an X () without sending it through (cleanedUp action) first. Fantastic!

---

I wish to improve the NeedsCleanup wrapper, so that it also "monadically" passes data, indicating what exactly needs cleaning up.

This is because, I have found, different NeedsCleanup X () actions may require different things to be cleaned up, and I have to clean up after all that have been binded together.

To be more precise, for each NeedsCleanup X t action, I would like for there to be associated a CleanupData:

data CleanupData = CleanupData
  {
       keyboard_needs_cleanup :: Bool
     , buttons_needing_cleanup :: Set.Set Buttons

     -- any other fields
     -- ...
  }

Two CleanupData can be combined, resulting in roughly a union ("afterwards, you have to clean up both for these actions").

-- | combines two CleanupData into the resulting CleanupData
combineCleanupData :: CleanupData -> CleanupData -> CleanupData
combineCleanupData dta1 dta2 =
  CleanupData
    {
         keyboard_needs_cleanup =
           (keyboard_needs_cleanup dta1) || (keyboard_needs_cleanup dta2)

       , buttons_needing_cleanup =
           (buttons_needing_cleanup dta1) `Set.union` (buttons_needing_cleanup dta2)

      -- union other data fields
      -- ...
    }

For example, if:

action1 :: NeedsCleanup X () is associated with dta1 :: CleanupData

action2 :: NeedsCleanup X () is associated with dta2 :: CleanupData

Then, action1 >> action2 should be associated with combineCleanupData dta1 dta2 (roughly "what you need to clean up for both").

Finally, at the end, the function cleanedUp :: NeedsCleanup X t -> X t should execute the underlying X t action and get the action's CleanupData (to see what needs cleaning up).

Is it possible to use a monadic wrapper to keep track of data in this way?

---

Update:

I ended up using something similar to Ilmo Euro's answer, except defining a Monoid structure for CleanupData instead of using the List Monoid. Something similar to:

import Control.Monad.Writer.Lazy (WriterT(..), runWriterT, tell, MonadWriter(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Monoid (Monoid(..))

initialCleanupData =
  CleanupData
    {
        keyboard_needs_cleanup = False
      , buttons_needing_cleanup = Set.empty

      -- initial values for other fields
    }

instance Monoid CleanupData where
  mempty = initialCleanupData
  mappend = combineCleanupData

newtype NeedsCleanup m t = 
  NeedsCleanup
    {
      to_writable :: WriterT CleanupData m t
    } deriving (MonadTrans, Monad, Applicative, Functor, MonadIO, MonadWriter CleanupData)

cleanup :: NeedsCleanup X t -> X t
cleanup action = do
  (ret_val, cleanup_data) <- runWriterT (to_writable action)

  -- clean up based on cleanup_data
  --   ...

  return ret_val 

In order to define an action that needs cleanup, I would tell it its CleanupData, for example, something similar to:

needsCleanup_GrabButton
  :: MonadIO m => Display -> Window -> Button -> NeedsCleanup m ()
needsCleanup_GrabButton dply window button = do
    liftIO $ grabButton dply button anyModifier window True buttonReleaseMask grabModeAsync grabModeAsync none none

    tell cleanup_data
  where
    -- the stuff we need to clean up from this
    -- particular action
    cleanup_data = initialCleanupData
      {
          buttons_needing_cleanup = Set.singleton button
      }

Answer 1

You can use, for example, the Writer monad for that:

import Control.Monad.Writer

data DirtyThing = Keyboard | Mouse
newtype Dirty m a = Dirty { unDirty :: WriterT [DirtyThing] m a }

doFoo :: Dirty IO ()
doFoo = -- doing something dirty

cleanup :: Dirty m a -> m a
cleanup action = do
    (val, dirtyThings) <- runWriterT (unDirty action)
    -- cleanup dirtyThings
    return val

For efficiency, you could use Set instead of lists (and define a newtype wrapper for it with an appropriate Monoid instance). Another, more type-safe (but much more tedious) way would be to use indexed monads.