SumBag: Add API documentation

src/Data/SumBag.hs

 {-# LANGUAGE RoleAnnotations #-}
 
+-- | == Multisets with constant-time fold.
+--
+-- This implements a generic (multi-)set data structure for __commutative__
+-- monoids with O(1) fold over the elements. Similar restrictions as for
+-- "Data.Set" apply. In particular, the number of /different/ elements in the
+-- Multiset must not exceed @maxBound::Int@. If this condition is violated, the
+-- behaviour is undefined.
+--
+-- == Implementation
+--
+-- Internally, a running total is kept and updated each time an element is
+-- inserted or deleted. The implementation is derivided from
+--
+--   * Stephen Adams, \"/Efficient sets: a balancing act/\",
+--     Journal of Functional Programming 3(4):553-562, October 1993,
+--     <http://www.swiss.ai.mit.edu/~adams/BB/>,
+--
+-- with the addition of the monoidal running total.
 module Data.SumBag
   ( SumBag
   , empty
@@ -16,22 +34,21 @@ import Prelude hiding (sum, min, elem)
 import Data.Foldable hiding (sum,elem)
 import qualified Data.List.NonEmpty as NE
 
-type SumBag a = Tree a
-
-data Tree a = Leaf | Node (MetaData a) (Element a) (Tree a) (Tree a)
+-- | A multiset of value @a@.
+data SumBag a = Leaf | Node (MetaData a) (Element a) (SumBag a) (SumBag a)
   deriving (Show)
 
-type role Tree nominal
+type role SumBag nominal
 
-instance (Ord a, Eq a) => Eq (Tree a) where
+instance (Ord a, Eq a) => Eq (SumBag a) where
   x == y = toAscList x == toAscList y
 
--- TODO There are a few functions from foldable that can be implemented way more
+-- TODO There are a few functions from Foldable that can be implemented way more
 -- efficiently.
 --
 -- Notably 'minimum' and 'maximum', but also explicit recursion instead of
 -- conversions to lists in a lot of cases.
-instance Foldable Tree where
+instance Foldable SumBag where
   foldMap f = foldMap f . toAscList
   {-# INLINE foldMap #-}
 
@@ -55,21 +72,39 @@ data Element a = Element
   deriving (Show)
 
 
+-- | The empty set.
+--
+-- Running time: O(1)
 empty :: SumBag a
 empty = Leaf
 
+-- | Constructs a set with a single element @a@.
 singleton :: Monoid a => a -> SumBag a
 singleton a =
   node (Element a (NE.fromList [a])) Leaf Leaf
 
+-- | Returns the number of nodes in the internal tree. This doesn't count
+-- duplicate elements and only returns the size of the internal tree.
+--
+-- Running time: O(1)
 size :: SumBag a -> Int
 size Leaf = 0
 size (Node node _ _ _) = nodeSize node
 
+-- | Compute the sum of all elements with their '<>' implementation. This is
+-- also called 'mconcat' and 'fold' for other containers.
+--
+-- Note that for the implementation to work, the Monoid @a@ has to be
+-- commutative.
+--
+-- Running time: O(1), since this value is cached internally.
 sum :: Monoid a => SumBag a -> a
 sum Leaf = mempty
 sum (Node node _ _ _) = nodeSum node
 
+-- | Inserts an element into the set.
+--
+-- Running time: O(log n)
 insert :: (Ord a, Monoid a) => a -> SumBag a -> SumBag a
 insert a Leaf = node (Element a (NE.fromList [a])) Leaf Leaf
 insert a (Node _ e left right)
@@ -77,6 +112,9 @@ insert a (Node _ e left right)
   | a > value e = balance1 e left (insert a right)
   | otherwise = node (addOnce e) left right
 
+-- | Tests membership of an element.
+--
+-- Running time: O(log n)
 elem :: (Ord a) => a -> SumBag a -> Bool
 elem _ Leaf = False
 elem a (Node _ e left right)
@@ -84,6 +122,10 @@ elem a (Node _ e left right)
   | a > value e = elem a right
   | otherwise = True
 
+-- | Delete an element from the set. If this element is not present, the
+-- original set is returned unmodified.
+--
+-- Running time: O(log n)
 delete :: (Ord a, Monoid a) => a -> SumBag a -> SumBag a
 delete _ Leaf = Leaf
 delete a (Node _ e left right)
@@ -98,6 +140,10 @@ delete a (Node _ e left right)
           let (min, rest) = delmin right
           in balance1 min left rest
 
+-- | Returns a list of all elements in the set in ascending order. Duplicate
+-- elements are correctly returned multiple times.
+--
+-- Running time: O(n)
 toAscList :: SumBag a -> [a]
 toAscList bag = helper bag []
   where helper Leaf accu = accu
@@ -106,13 +152,16 @@ toAscList bag = helper bag []
 
         mkList (Element val mult) = map (const val) (NE.toList mult)
 
+-- | Return a multiset of all elements of a list.
+--
+-- Running time: O(n · log n)
 fromList :: (Ord a, Monoid a) => [a] -> SumBag a
 fromList = foldr insert empty
 
 -- Internal functions
 
 -- | "Smart" constructor for Node. Will compute the meta data from its subtrees
-node :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+node :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 node a left right =
   let nodeData = MetaData
         { nodeSize = size left + 1 + size right
@@ -120,24 +169,24 @@ node a left right =
         }
   in Node nodeData a left right
 
-rotateSingleLeft :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+rotateSingleLeft :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 rotateSingleLeft a x (Node _ b y z) = node b (node a x y) z
 rotateSingleLeft _ _ _ = error "rotateSingleLeft called with empty right tree"
 
-rotateSingleRight :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+rotateSingleRight :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 rotateSingleRight b (Node _ a x y) z = node a x (node b y z)
 rotateSingleRight _ _ _ = error "rotateSingleRight called with empty left tree"
 
-rotateDoubleLeft :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+rotateDoubleLeft :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 rotateDoubleLeft a x (Node _ c (Node _ b y1 y2) z) = node b (node a x y1) (node c y2 z)
 rotateDoubleLeft _ _ _ = error "rotateDoubleLeft called with too small right tree"
 
-rotateDoubleRight :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+rotateDoubleRight :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 rotateDoubleRight c (Node _ a x (Node _ b y1 y2)) z = node b (node a x y1) (node c y2 z)
 rotateDoubleRight _ _ _ = error "rotateDoubleRight called with too small left tree"
 
 
-balance1 :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
+balance1 :: Monoid a => Element a -> SumBag a -> SumBag a -> SumBag a
 balance1 a left right
   -- Subtrees have only one element
   | size left + size right < 2 = node a left right
@@ -165,7 +214,7 @@ delOnce e = case snd (NE.uncons (multiplicity e)) of
   Nothing -> Nothing
   Just rest -> Just (e { multiplicity = rest })
 
-delmin :: Monoid a => Tree a -> (Element a, Tree a)
+delmin :: Monoid a => SumBag a -> (Element a, SumBag a)
 delmin Leaf = error "delmin: Empty tree"
 delmin (Node _ e Leaf right) = (e, right)
 delmin (Node _ e left right) = (\left' -> balance1 e left' right) <$> delmin left