Destroy fancy typing, and some work

1 files changed, 14 insertions, 3 deletions

diff --git a/src/Data/Bag.hs b/src/Data/Bag.hs
index 9cdce89..f1173e4 100644
--- a/src/Data/Bag.hs
+++ b/src/Data/Bag.hs
@@ -3,12 +3,16 @@ module Data.Bag where
 
 
 data Bag a
-  = BTwo (Bag a) (Bag a)
+  = BZero
   | BOne a
-  | BZero
+  | BTwo (Bag a) (Bag a)
   deriving (Functor, Foldable, Traversable)
 
-instance Semigroup (Bag a) where (<>) = BTwo
+instance Semigroup (Bag a) where
+  BZero <> b = b
+  b <> BZero = b
+  b1 <> b2 = BTwo b1 b2
+
 instance Monoid (Bag a) where mempty = BZero
 
 instance Applicative Bag where
@@ -18,3 +22,10 @@ instance Applicative Bag where
   _ <*> BZero = BZero
   BOne f <*> b = f <$> b
   BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
+
+bagPartition :: (a -> Maybe b) -> Bag a -> (Bag b, Bag a)
+bagPartition _ BZero = (BZero, BZero)
+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