monoid-comprehensions.hs

#! /usr/bin/env nix-shell
#! nix-shell -I nixpkgs=channel:nixos-18.09 -i ghci -p "haskellPackages.ghcWithPackages (p: [p.monoidal-containers p.average])"

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MonadComprehensions #-}

import           Control.Applicative ( Alternative(..) )
import           Control.Monad.Trans.Cont
import           Data.Map.Monoidal ( MonoidalMap )
import qualified Data.Map.Monoidal as Map
import           Data.Monoid ( Sum(..) )
import           Data.Monoid.Average
import           Data.Ord ( Down(..) )
import qualified Data.Set as Set

newtype MonoidComp r a = MonoidComp { unMonoidComp :: Cont r a }
  deriving (Functor, Applicative, Monad)

instance Monoid r => Alternative (MonoidComp r) where
  empty = MonoidComp $ cont $ const mempty
  MonoidComp a <|> MonoidComp b =
    MonoidComp $ cont $ \f -> runCont a f <> runCont b f

fromFoldable :: (Foldable f, Monoid r) => f a -> MonoidComp r a
fromFoldable = MonoidComp . cont . flip foldMap

runMonoidComp :: MonoidComp r r -> r
runMonoidComp = evalCont . unMonoidComp

ex1
  :: (Monoid c, Foldable f1, Foldable f2)
  => (a -> b -> Bool)
  -> (a -> b -> c)
  -> f1 a
  -> f2 b
  -> c
ex1 f g xs ys =
  runMonoidComp [ g x y | x <- fromFoldable xs, y <- fromFoldable ys, f x y ]

-- ghci> ex2
-- Sum {getSum = 36}
ex2 :: Sum Int
ex2 = runMonoidComp
  [ Sum (length y + x)
  | x <- fromFoldable [1, 2, 3]
  , y <- fromFoldable $ Set.fromList ["a", "bb", "ccc"]
  ]

groupBy :: Ord k => a -> k -> MonoidalMap k a
groupBy = flip Map.singleton

data Person = Person { salary :: Float, age :: Int }

ex3 :: Foldable f => f Person -> MonoidalMap Int (Average Float)
ex3 persons = runMonoidComp
  [ Average [salary p] `groupBy` age p | p <- fromFoldable persons ]


ex4
  :: Foldable f
  => f (Int, Int)
  -> (MonoidalMap Int (Sum Int), MonoidalMap Int (Average Int))
ex4 ls = runMonoidComp
  [ (Sum x `groupBy` y, Average [y] `groupBy` x) | (x, y) <- fromFoldable ls ]

ex5 :: (Monoid a, Foldable t, Num a) => [(t a, Int)] -> [a]
ex5 xs = take 2 $ Map.elems $ runMonoidComp
  [ sum x `groupBy` y | (x, y) <- fromFoldable (drop 1 xs) ]

orderBy :: Ord k => a -> k -> MonoidalMap k a
orderBy = flip Map.singleton

ex6
  :: (Foldable f, Ord z, Ord y)
  => f (x, y, z)
  -> MonoidalMap z (Sum Int, MonoidalMap (Down y) [x])
ex6 xs = runMonoidComp
  [ (count, [x] `orderBy` Down y) `groupBy` z | (x, y, z) <- fromFoldable xs ]
  where count = Sum 1