{-# LANGUAGE DeriveTraversable #-}
module Data.Bag where


data Bag a
  = BZero
  | BOne a
  | BTwo (Bag a) (Bag a)
  deriving (Functor, Foldable, Traversable)

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
  pure = BOne

  BZero <*> _ = BZero
  _ <*> 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