@@ -73,6 +73,7 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
7373 , foldFilter
7474 , deleteBy
7575 , uniq
76+ , nub
7677
7778 -- * Trimming
7879 -- | Produce a subset of the stream trimmed at ends.
@@ -135,6 +136,7 @@ import Data.Maybe (fromJust, isJust)
135136import Fusion.Plugin.Types (Fuse(..))
136137import GHC.Types (SPEC(..))
137138import qualified Control.Monad.Catch as MC
139+import qualified Data.Set as Set
138140
139141import Streamly.Internal.Control.Concurrent (MonadAsync)
140142import Streamly.Internal.Control.ForkLifted (fork, forkManaged)
@@ -796,6 +798,26 @@ uniq (Stream step state) = Stream step' (Nothing, state)
796798 Skip s -> return $ Skip (Just x, s)
797799 Stop -> return Stop
798800
801+-- | The memory used is proportional to the number of unique elements in the
802+-- stream. If we want to limit the memory we can just use "take" to limit the
803+-- uniq elements in the stream.
804+{-# INLINE_NORMAL nub #-}
805+nub :: (Monad m, Ord a) => Stream m a -> Stream m a
806+nub (Stream step1 state1) = Stream step (Set.empty, state1)
807+
808+ where
809+
810+ step gst (set, st) = do
811+ r <- step1 gst st
812+ return
813+ $ case r of
814+ Yield x s ->
815+ if Set.member x set
816+ then Skip (set, s)
817+ else Yield x (Set.insert x set, s)
818+ Skip s -> Skip (set, s)
819+ Stop -> Stop
820+
799821{-# INLINE_NORMAL deleteBy #-}
800822deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
801823deleteBy eq x (Stream step state) = Stream step' (state, False)