Partition.hs

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

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 From :: k -> Type
data From k where
  Id :: From Type
  TyCon :: (k -> Type) -> From k

type ($) :: From k -> k -> Type
type family f $ t where
  Id $ t = t
  TyCon f $ t = f t

infixr 0 $

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

type Product :: From k -> [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 (TyCon f) ts

instance Countable '[] where
  duplicate _ = Nil

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

partition :: Countable ts => (t -> Sum Id ts) -> [t] -> Product (TyCon []) ts
partition = \f -> foldl (flip (prepend . f)) (duplicate []) where
  prepend :: Sum Id ts -> Product (TyCon []) ts -> Product (TyCon []) ts
  prepend = \case
    This a -> \(as :* p) -> (a : as) :* p
    That b -> \(as :* p) -> as :* prepend b p

-- quality of life utilities

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

pattern Case1 :: f $ b -> Sum f (a ': b ': ts)
pattern Case1 b = That (This b)

pattern Case2 :: f $ c -> Sum f (a ': b ': c ': ts)
pattern Case2 c = That (That (This c))

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