SetDiff Monoid and group

module SetDiff
    ( SetDiff
    , fromList
    , fromSet
    , apply
    , toSet
    , inverse
    ) where

import Data.Monoid
import Data.List
import Data.Set (Set, (\\))
import qualified Data.Set as Set
import Data.Map (Map)
import qualified Data.Map as Map


-- | A 'SetDiff' represents a set of additions and removals of
-- elements from any 'Set' of the same element type.
data SetDiff a = SetDiff (Map a Int) deriving Show

-- | Make a 'SetDiff' from a list of elements.  The elements are
-- interpreted as additions.
fromList :: Ord a => [a] -> SetDiff a
fromList xs = SetDiff (Map.fromList (map withOne xs))
    where withOne x = (x, 1)

-- | Same as 'fromList' but takes a 'Set' instead.
fromSet :: Ord a => Set a -> SetDiff a
fromSet = fromList . Set.toList

-- | Apply a 'SetDiff' to a 'Set'.  Positive elements of the 'SetDiff'
-- will be added to the 'Set', negative ones will be removed.  Zeros
-- will neither be added nor removed.
apply :: Ord a => SetDiff a -> Set a -> Set a
apply (SetDiff m) xs = (xs `Set.difference` neg) `Set.union` pos
    where (neg', nonneg') = partition ((<0) . snd) (Map.toList m)
          neg = Set.fromList (map fst neg')
          pos = Set.fromList (map fst (filter ((>0) . snd) nonneg'))

-- | Turn a 'SetDiff' into a 'Set'.  This is the same as applying the
-- 'SetDiff' to an empty 'Set'.
toSet :: Ord a => SetDiff a -> Set a
toSet diff = apply diff mempty

-- | 'SetDiff' is a 'Monoid'; you can combine two 'SetDiff's into a
-- composite one that reflects the result of applying both.  Laws:
--
-- > apply mempty xs == xs
-- > apply diff1 (apply diff2 xs) == apply (diff1 `mappend` diff2) xs
instance Ord a => Monoid (SetDiff a) where
    mempty = SetDiff Map.empty
    (SetDiff m0) `mappend` (SetDiff m1) =
        SetDiff (Map.unionWith (+) m0 m1)

-- | 'SetDiff' is a group.  Law:
--
-- > diff `mappend` inverse diff == mempty
inverse :: SetDiff a -> SetDiff a
inverse (SetDiff m) = SetDiff (Map.map negate m)