Partition.hs

{-# OPTIONS_GHC -Wall -Wextra -Werror #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
module Partition where

import Data.Functor.Identity
import Data.Kind

data Animal = Cat | Rabbit | Snake | Dog | Pig | Sheep | Wildebeest
data Monster = Dracula | Wolfman | Vampire | Zombie | Kaiju
data Meal = Breakfast | Brunch | Lunch | Tea | Dinner | Picnic
data Thing = Animal Animal | Monster Monster | Meal Meal

example :: [Thing] -> ([Animal], [Monster], [Meal])
example = toTriple . partition \case
  Animal animal -> Case0 animal
  Monster monster -> Case1 monster
  Meal meal -> Case2 meal

----

type Sum :: (k -> Type) -> [k] -> Type
data Sum f ts where
  This :: f t -> Sum f (t ': ts)
  That :: Sum f ts -> Sum f (t ': ts)

type Product :: (k -> Type) -> [k] -> Type
data Product f ts where
  Nil :: Product f '[]
  (:*) :: f t -> Product f ts -> Product f (t ': ts)

infixr 4 :*

type Countable :: [Type] -> Constraint
class Countable ts where
  duplicate :: (forall t. f t) -> Product f ts

instance Countable '[] where
  duplicate _ = Nil

instance Countable ts => Countable (t ': ts) where
  duplicate ft = ft :* duplicate ft

partition :: Countable ts => (t -> Sum Identity ts) -> [t] -> Product [] ts
partition = \f -> foldl (flip (prepend . f)) (duplicate []) where
  prepend :: Sum Identity ts -> Product [] ts -> Product [] ts
  prepend = \case
    -- we could get rid of the superfluous Identity wrapper with more type-fu
    -- https://gist.github.com/rampion/84d39cdbe337ebe8b6b5dda07b5c3774
    This (Identity a) -> \(as :* p) -> (a : as) :* p
    That b -> \(as :* p) -> as :* prepend b p

-- quality of life utilities

pattern Case0 :: a -> Sum Identity (a ': ts)
pattern Case0 a = This (Identity a)

pattern Case1 :: b -> Sum Identity (a ': b ': ts)
pattern Case1 b = That (Case0 b)

pattern Case2 :: c -> Sum Identity (a ': b ': c ': ts)
pattern Case2 c = That (Case1 c)

toTriple :: Product f [a,b,c] -> (f a, f b, f c)
toTriple (fa :* fb :* fc :* Nil) = (fa, fb, fc)