1 files changed, 62 insertions, 3 deletions
diff --git a/src/Data/Bag.hs b/src/Data/Bag.hs
index f1173e4..ba1f912 100644
--- a/src/Data/Bag.hs
+++ b/src/Data/Bag.hs
@@ -1,19 +1,41 @@
 {-# LANGUAGE DeriveTraversable #-}
-module Data.Bag where
+module Data.Bag (
+  Bag,
+  bagFromList,
+  bagFilter,
+  bagPartition,
+) where
+
+import Data.Bifunctor (bimap)
+import Data.Foldable (toList)
+import Data.List.NonEmpty (NonEmpty((:|)))
+import Data.Maybe (mapMaybe)
+import Data.Semigroup
 
 
 data Bag a
   = BZero
   | BOne a
   | BTwo (Bag a) (Bag a)
-  deriving (Functor, Foldable, Traversable)
+  | BList [Bag a]  -- make mconcat efficient
+  | BList' [a]  -- make bagFromList efficient
+  deriving (Functor, Traversable)
+
+instance Show a => Show (Bag a) where
+  showsPrec d b = showParen (d > 10) $
+    showString "Bag " . showList (toList b)
 
 instance Semigroup (Bag a) where
   BZero <> b = b
   b <> BZero = b
   b1 <> b2 = BTwo b1 b2
 
-instance Monoid (Bag a) where mempty = BZero
+  sconcat (b :| bs) = b <> mconcat bs
+  stimes n b = mconcat (replicate (fromIntegral n) b)
+
+instance Monoid (Bag a) where
+  mempty = BZero
+  mconcat = BList
 
 instance Applicative Bag where
   pure = BOne
@@ -22,6 +44,36 @@ instance Applicative Bag where
   _ <*> BZero = BZero
   BOne f <*> b = f <$> b
   BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
+  BList bs <*> b = BList (map (<*> b) bs)
+  BList' xs <*> b = BList (map BOne xs) <*> b
+
+instance Foldable Bag where
+  foldMap _ BZero = mempty
+  foldMap f (BOne x) = f x
+  foldMap f (BTwo b1 b2) = foldMap f b1 <> foldMap f b2
+  foldMap f (BList l) = foldMap (foldMap f) l
+  foldMap f (BList' l) = foldMap f l
+
+  toList (BList' xs) = xs
+  toList b = foldr (:) [] b
+
+  null BZero = True
+  null BOne{} = False
+  null (BTwo b1 b2) = null b1 && null b2
+  null (BList l) = all null l
+  null (BList' l) = null l
+
+bagFromList :: [a] -> Bag a
+bagFromList = BList'
+
+bagFilter :: (a -> Maybe b) -> Bag a -> Bag b
+bagFilter _ BZero = BZero
+bagFilter f (BOne x)
+  | Just y <- f x = BOne y
+  | otherwise = BZero
+bagFilter f (BTwo b1 b2) = bagFilter f b1 <> bagFilter f b2
+bagFilter f (BList bs) = BList (map (bagFilter f) bs)
+bagFilter f (BList' xs) = BList' (mapMaybe f xs)
 
 bagPartition :: (a -> Maybe b) -> Bag a -> (Bag b, Bag a)
 bagPartition _ BZero = (BZero, BZero)
@@ -29,3 +81,10 @@ bagPartition f (BOne x)
   | Just y <- f x = (BOne y, BZero)
   | otherwise = (BZero, BOne x)
 bagPartition f (BTwo b1 b2) = bagPartition f b1 <> bagPartition f b2
+bagPartition f (BList bs) = foldMap (bagPartition f) bs
+bagPartition f (BList' xs) =
+  bimap bagFromList bagFromList $
+    foldr (\x (l1,l2) -> case f x of
+                           Just y -> (y : l1, l2)
+                           Nothing -> (l1, x : l2))
+          ([], []) xs