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
  , singleton
  , sum
  , insert
  , delete
  , elem
  , toAscList
  , fromList
  ) where

import Prelude hiding (sum, min, elem)
import Data.Foldable hiding (sum,elem)
import qualified Data.List.NonEmpty as NE

-- | A multiset of value @a@.
data SumBag a = Leaf | Node (MetaData a) (Element a) (SumBag a) (SumBag a)
  deriving (Show)

type role SumBag nominal

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
-- efficiently.
--
-- Notably 'minimum' and 'maximum', but also explicit recursion instead of
-- conversions to lists in a lot of cases.
instance Foldable SumBag where
  foldMap f = foldMap f . toAscList
  {-# INLINE foldMap #-}

  fold = sum
  {-# INLINE fold #-}

  toList = toAscList
  {-# INLINE toList #-}
  

data MetaData a = MetaData
  { nodeSize :: Int
  , nodeSum :: a
  }
  deriving (Show)

data Element a = Element
  { value :: a
  , multiplicity :: NE.NonEmpty a
  }
  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)
  | a < value e = balance1 e (insert a 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)
  | a < value e = elem a left
  | 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)
  | a < value e = balance1 e (delete a left) right
  | a > value e = balance1 e left (delete a right)
  | Just e' <- delOnce e = node e' left right
  | otherwise = helper left right

  where helper Leaf right = right
        helper left Leaf = left
        helper 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
        helper (Node _ e left right) accu =
          helper left (mkList e ++ helper right accu)

        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 -> SumBag a -> SumBag a -> SumBag a
node a left right =
  let nodeData = MetaData
        { nodeSize = size left + 1 + size right
        , nodeSum = NE.head (multiplicity a) <> sum left <> sum right
        }
  in Node nodeData a left right

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 -> 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 -> 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 -> 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 -> SumBag a -> SumBag a -> SumBag a
balance1 a left right
  -- Subtrees have only one element
  | size left + size right < 2 = node a left right
  -- Right subtree is too heavy
  | size right > balanceBound * size left =
    let Node _ _ rleft rright = right
        sizeRL = size rleft
        sizeRR = size rright
    in if sizeRL < sizeRR then rotateSingleLeft a left right else rotateDoubleLeft a left right
  -- Left subtree is too heavy
  | size left > balanceBound * size right =
    let Node _ _ lleft lright = left
        sizeLL = size lleft
        sizeLR = size lright
    in if sizeLR < sizeLL then rotateSingleRight a left right else rotateDoubleRight a left right
  -- No subtree is too heavy, we can just form a new tree straight away
  | otherwise = node a left right

addOnce :: Semigroup a => Element a -> Element a
addOnce e = let total = NE.head (multiplicity e)
            in e { multiplicity = NE.cons (total <> value e) (multiplicity e) }

delOnce :: Element a -> Maybe (Element a)
delOnce e = case snd (NE.uncons (multiplicity e)) of
  Nothing -> Nothing
  Just rest -> Just (e { multiplicity = rest })

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

balanceBound :: Int
balanceBound = 4