module Mikan.Utils.Set
  ( lookupKey
    -- * Re-exports
  , module Data.Set
  ) where

import Data.Set
import Data.Set.Internal (Set(Tip, Bin))

-- Adapted from rejected pull request to containers
-- https://github.com/haskell/containers/pull/291

-- | \(\mathcal{O}(\log n)\). Find the given element and return the copy contained in the set.
lookupKey :: Ord a => a -> Set a -> Maybe a
lookupKey :: forall a. Ord a => a -> Set a -> Maybe a
lookupKey !a
x = Set a -> Maybe a
go
  where
    go :: Set a -> Maybe a
go Set a
Tip = Maybe a
forall a. Maybe a
Nothing
    go (Bin Size
_ a
y Set a
l Set a
r) = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
      Ordering
LT -> Set a -> Maybe a
go Set a
l
      Ordering
GT -> Set a -> Maybe a
go Set a
r
      Ordering
EQ -> a -> Maybe a
forall a. a -> Maybe a
Just a
y