@@ -55,6 +55,9 @@ module Streamly.Internal.Data.Fold
5555 , drainBy
5656 , last
5757 , length
58+ , genericLength
59+ , countDistinct
60+ , countDistinctInt
5861 , mean
5962 , variance
6063 , stdDev
@@ -268,6 +271,7 @@ import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
268271import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
269272import Streamly.Internal.Data.Stream.Serial (SerialT(..))
270273
274+import qualified Data.IntSet as IntSet
271275import qualified Data.Map.Strict as Map
272276import qualified Data.Set as Set
273277import qualified Streamly.Internal.Data.Pipe.Type as Pipe
@@ -526,6 +530,56 @@ genericLength = foldl' (\n _ -> n + 1) 0
526530length :: Monad m => Fold m a Int
527531length = genericLength
528532
533+-- XXX Try Hash set
534+-- XXX Add a countDistinct window fold
535+-- XXX Add a bloom filter fold
536+
537+-- | Count non-duplicate elements in the stream.
538+--
539+-- The memory used is proportional to the number of distinct elements in the
540+-- stream, to guard against using too much memory use it as a scan and
541+-- terminate if the count reaches more than a threshold.
542+--
543+-- /Space/: \(\mathcal{O}(n)\)
544+--
545+-- /Pre-release/
546+--
547+{-# INLINE countDistinct #-}
548+countDistinct :: (Monad m, Ord a) => Fold m a Int
549+countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
550+
551+ where
552+
553+ initial = Tuple' Set.empty 0
554+
555+ step (Tuple' set n) x = do
556+ if Set.member x set
557+ then
558+ Tuple' set n
559+ else
560+ let cnt = n + 1
561+ in Tuple' (Set.insert x set) cnt
562+
563+-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
564+-- performance.
565+--
566+-- /Pre-release/
567+{-# INLINE countDistinctInt #-}
568+countDistinctInt :: Monad m => Fold m Int Int
569+countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
570+
571+ where
572+
573+ initial = Tuple' IntSet.empty 0
574+
575+ step (Tuple' set n) x = do
576+ if IntSet.member x set
577+ then
578+ Tuple' set n
579+ else
580+ let cnt = n + 1
581+ in Tuple' (IntSet.insert x set) cnt
582+
529583-- | Determine the sum of all elements of a stream of numbers. Returns additive
530584-- identity (@0@) when the stream is empty. Note that this is not numerically
531585-- stable for floating point numbers.