Ugly Folds

README.md

This is a strict, somewhat efficient implementation of something between Data.List.partition and Beautiful Folding.

It allows a (pure or ST) stream of values to be processed into multiple fields, using a nice 'Applicative' interface.

Beautiful Folding allocates O(#fields * #inputs), whereas this allocates O(#fields + #inputs).

Time complexity may still be O(#fields * #inputs), because of the nested tuple reads, but presumably with a lower constant factor than Beautiful Folding. It would take a clever compiler to see that the m type can be represented by a flat data structure instead of a heterogeneous linked list composed of 2-tuples.

Unlike partition, it does not guarantee that a value only occurs in one field; it is a generalization after all. Partitioner should probably be renamed because of this.

I wrote this module out of curiosity as a sort of feasibility study in case I had to optimize a pure implementation of this concept, which has the same interface but just looks a lot like a simpler Beautiful Folding on the inside.

Partitioner.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}


-- | This is a strict, somewhat efficient implementation of something between 'Data.List.partition' and Beautiful Folding.
--
-- It allows a (pure or ST) stream of values to be processed into multiple fields, using a nice 'Applicative' interface.
--
-- Beautiful Folding allocates O(#fields * #inputs), whereas this allocates O(#fields + #inputs).
--
-- Time complexity may still be O(#fields * #inputs), because of the nested tuple reads, but presumably with a lower constant factor than Beautiful Folding. It would take a clever compiler to see that the @m@ type can be represented by a flat data structure instead of a heterogeneous linked list composed of 2-tuples.
--
-- Unlike 'Data.List.partition', it does not guarantee that a value only occurs in one field; it is a generalization after all.

module Data.Functor.Partitioner where

import Control.Monad.ST (ST, runST)
import Data.Foldable (for_, toList)
import Data.Monoid (Ap (Ap))
import Data.STRef (modifySTRef, newSTRef, readSTRef)
import Prelude

newtype Partitioner a b = Partitioner (forall s. PartitionerST s a b)
  deriving (Functor)

data PartitionerST s a b = forall m.
  PartitionerST
  { pstInit :: ST s m,
    pstIteration :: a -> m -> ST s (),
    pstFinish :: m -> ST s b
  }

-- | Based on the 'Applicative', via 'Ap'
deriving via (Ap (Partitioner a) b) instance Semigroup b => Semigroup (Partitioner a b)

-- | Based on the 'Applicative', via 'Ap'
deriving via (Ap (Partitioner a) b) instance Monoid b => Monoid (Partitioner a b)

-- | Based on the 'Applicative', via 'Ap'
deriving via (Ap (Partitioner a) b) instance Num b => Num (Partitioner a b)

-- | Based on the 'Applicative', via 'Ap'
deriving via (Ap (Partitioner a) b) instance Eq b => Eq (Partitioner a b)

-- | Based on the 'Applicative', via 'Ap'
deriving via (Ap (Partitioner a) b) instance Ord b => Ord (Partitioner a b)

-- | Core interface
instance Applicative (Partitioner a) where
  pure a = Partitioner (pure a)
  Partitioner p <*> Partitioner q = Partitioner (p <*> q)

partitionList :: Partitioner a b -> [a] -> b
partitionList p as = runST $ case p of
  Partitioner PartitionerST {pstInit = pinit, pstIteration = piter, pstFinish = pfin} -> do
    m <- pinit
    for_ as $ \a ->
      piter a m
    pfin m

part :: Foldable f => (a -> f b) -> Partitioner a [b]
part f = parts (\a -> toList (f a))

parts :: (a -> [b]) -> Partitioner a [b]
parts f =
  ($ []) <$> Data.Functor.Partitioner.foldl (\bs a -> bs . (f a ++)) id

foldl :: (b -> a -> b) -> b -> Partitioner a b
foldl op nul =
  Partitioner
    PartitionerST
      { pstInit = newSTRef nul,
        pstIteration = \a ref -> modifySTRef ref (\b -> op b a),
        pstFinish = \b -> readSTRef b
      }

instance Functor (PartitionerST s a) where
  fmap f (PartitionerST pinit piter pfin) =
    PartitionerST
      { pstInit = pinit,
        pstIteration = piter,
        pstFinish = fmap f . pfin
      }

instance Applicative (PartitionerST s a) where
  pure a = PartitionerST {pstInit = pure (), pstIteration = mempty, pstFinish = const (pure a)}

  PartitionerST {pstInit = pinit, pstIteration = piter, pstFinish = pfin}
    <*> PartitionerST {pstInit = qinit, pstIteration = qiter, pstFinish = qfin} =
      PartitionerST
        { pstInit = (,) <$> pinit <*> qinit,
          pstIteration = \a (mi, mj) -> piter a mi *> qiter a mj,
          pstFinish = \(mi, mj) -> pfin mi <*> qfin mj
        }

-- instance Profunctor Partitioner
-- ...