Haskell Write Monad for expressions

I am trying to design embedded language, where operations can raise certain flags depending on values. I foresee operation on scalar values as well as on vectors (e.g. map, fold, etc.) My idea is to use Writer Monad to keep track of flags. Simplified example, where actual type is "Int" and flag is raised if any of argument is 0.

import Control.Monad.Identity
import Control.Monad.Writer
import Data.Monoid    

type WInt = Writer Any Int

bplus :: Int -> Int -> WInt
bplus a b =
    do
      tell (Any (a == 0 || b == 0)) ;
           return (a+b)

wbplus :: WInt -> WInt -> WInt
wbplus wa wb =
    do
      a <- wa ;
      b <- wb ;
      tell (Any (a == 0 || b == 0)) ;
           return (a+b)

ex0 = runWriter (bplus 1 2) 
ex1 = runWriter (bplus 0 2)

ex2 = runWriter (wbplus (return 1) (return 2))
ex3 = runWriter (wbplus (return 0) (return 2))

ex4 = runWriter (wbplus (wbplus (return 1) (return 2)) (return 2))
ex5 = runWriter (wbplus (wbplus (return 0) (return 2)) (return 2))
ex6 = runWriter (wbplus (wbplus (return 1) (return 2)) (return 0))

I am little unsure what is the best way to implement this. Some questions:

Should I define all operations like I did for bplus or like for wbplus. Laters makes composition easier, it seems. But to use foldM binary operator should have type Int -> Int -> WInt.
What would be the appropriate type for lists: Writer Any [Int] or [Wint]?

Any suggestions or thoughts are appreciated.

Solution

You can derive bplus from wbplus and vice versa using the appropriate monadic operations:

import Control.Monad

apM2 :: Monad m => (a -> b -> m c) -> m a -> m b -> m c
apM2 f ma mb = do
  a <- ma
  b <- mb
  f a b

pureM2 :: Monad m => (m a -> m b -> m c) -> a -> b -> m c
pureM2 f a b = f (return a) (return b)

They are inverses of each other, evident from the type signatures of their compositions:

ghci> :t pureM2 . apM2
pureM2 . apM2 :: Monad m => (a -> b -> m c) -> a -> b -> m c

ghci> :t apM2 . pureM2
apM2 . pureM2 :: Monad m => (m a -> m b -> m c) -> m a -> m b -> m c

Now you can define wbplus = apM2 bplus or bplus = pureM2 wbplus. There's no definite answer which one is better, use your taste and judgement. TemplateHaskell goes with the wbplus approach and defines all operations to work with values in the Q monad. See Language.Haskell.TH.Lib.

Regarding [m a] vs m [a], you can only go in one direction (via sequence :: Monad m => [m a] -> m [a]). Would you ever want to go in the opposite direction? Do you care about individual values having their own flags or would you rather annotate the computation as a whole with flags?

The real question is, what is your mental model for this? However, let's think about some consequences of each design choice.

If you choose to represent each value as Writer Any a and have all operations work with it, you can start with a newtype:

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

import Control.Monad.Writer

newtype Value a = Value (Writer Any a)
  deriving (Functor, Applicative, Monad)

Now you can define instances of standard type classes for your values:

instance (Num a, Eq a) => Num (Value a) where
  va + vb = do
    a <- va
    b <- vb
    (Value . tell . Any) (b == 0 || a == 0)
    return (a + b)
  (*) = liftM2 (*)
  abs = fmap abs
  signum = fmap signum
  negate = fmap negate
  fromInteger = return . fromInteger

  instance Monoid a => Monoid (Value a) where
    mempty = pure mempty
    mappend = liftM2 mappend

For an EDSL this gives a huge advantage: terseness and syntactic support from the compiler. You can now write getValue (42 + 0) instead of wbplus (pure 42) (pure 0).

If, instead, you don't think about flags as a part of your values and rather see them as an external effect, it's better to go with the alternative approach. But rather than write something like Writer Any [Int], use corresponding classes from mtl: MonadWriter Any m => m [Int]. This way, if you later find out that you need to use other effects, you can easily add them to some (but not all) operations. For example, you might want to raise an error in case of division by zero:
```
  data DivisionByZero = DivisionByZero

  divZ :: (MonadError DivisionByZero m, Fractional a, Eq a) => a -> a -> m a
  divZ a b
    | b == 0 = throwError DivisionByZero
    | otherwise = pure (a / b)

  plusF :: (MonadWriter Any m, Num a, Eq a) => a -> a -> m a
  plusF a b = do
    tell (Any (b == 0 || a == 0))
    return (a + b)
```
Now you can use plusF and divZ together within one monad, although they have different effects. If you'll later find yourself in need to integrate with some external library, this flexibility will come in handy.

Now, I didn't give it too much thought, but perhaps you could combine those approaches using something like newtype Value m a = Value { getValue :: m a }. Good luck exploring the design space :)