Generalized partitionEithers

PartitionNS.hs

#!/usr/bin/env cabal
{- cabal:
build-depends: base       >= 4.9
             , QuickCheck >= 2.12
             , sop-core   >= 0.4
-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
module Main where

import Data.Coerce (coerce)
import Data.Either (partitionEithers)
import Data.SOP
import Test.QuickCheck (quickCheck)

type Listed f = [] :.: f

partitionNS :: SListI xs => [NS f xs] -> NP (Listed f) xs
partitionNS = foldr insertIntoNP emptyListNP
  where
    emptyListNP :: SListI xs => NP (Listed f) xs
    emptyListNP = para_SList Nil (Comp [] :*)

    insertIntoNP :: NS f xs -> NP (Listed f) xs -> NP (Listed f) xs
    insertIntoNP (Z fx) (Comp fxs :* fxss) = Comp (fx:fxs) :* fxss
    insertIntoNP (S ns) (np       :* nps)  = np :* insertIntoNP ns nps

partitionEithers' :: [Either a b] -> ([a], [b])
partitionEithers' = fromNP . partitionNS . map toNS
  where
    toNS :: Either a b -> NS I '[a, b]
    toNS (Left l)  = Z (I l)
    toNS (Right r) = S (Z (I r))

    fromNP :: NP (Listed I) '[a, b] -> ([a], [b])
    fromNP (as :* bs :* Nil) = coerce (as, bs)

main :: IO ()
main = quickCheck $ \(e :: [Either Int Char]) ->
                    partitionEithers e == partitionEithers' e