content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  -- | A tagged intset to keep the type of keys clear.  Mainly intended for
  -- Glean predicates.
  module Glean.Util.PredSet
    ( PredSet(..), fromList, fromDistinctAscList
    , empty, insert, toList, toAscList, size
    , member, notMember
    , singleton, union, difference
      -- * PredSet to/from PredMap
    , keysSet, fromSet, fromSetM
    ) where
  
  import Control.DeepSeq (NFData)
  import qualified Data.IntMap.Strict as IntMap
  import Data.IntSet (IntSet)
  import qualified Data.IntSet as IntSet
  import Prelude hiding (lookup)
  
  import Glean
  import Glean.Util.PredMap (PredMap)
  import qualified Glean.Util.PredMap as PredMap
  
  -- -----------------------------------------------------------------------------
  
  {-# INLINE one #-}
  one :: IdOf p -> Int
  one = fromIntegral . fromFid . idOf
  
  -- -----------------------------------------------------------------------------
  
  -- | I want some more type safety around IntSet. The 'p' will usually
  -- be a Predicate type, e.g. (IdOf p) as the logical key.
  newtype PredSet p = PredSet { predSet :: IntSet }
    deriving (Show, Semigroup, Monoid, NFData)
  
  fromList :: [IdOf p] -> PredSet p
  fromList = PredSet . IntSet.fromList . map one
  
  singleton :: IdOf p -> PredSet p
  singleton = PredSet . IntSet.singleton . one
  
  fromDistinctAscList :: [IdOf p] -> PredSet p
  fromDistinctAscList = PredSet . IntSet.fromDistinctAscList . map one
  
  empty :: PredSet p
  empty = PredSet IntSet.empty
  
  insert :: IdOf p -> PredSet p -> PredSet p
  insert p = PredSet . IntSet.insert (one p) . predSet
  
  -- | Synonym for 'toAscList'
  toList :: PredSet p -> [IdOf p]
  toList = toAscList
  
  toAscList :: PredSet p -> [IdOf p]
  toAscList = map (IdOf . Fid . fromIntegral) . IntSet.toAscList . predSet
  
  size :: PredSet p -> Int
  size = IntSet.size . predSet
  
  member :: IdOf p -> PredSet p -> Bool
  member k pm = IntSet.member (one k) (predSet pm)
  
  notMember :: IdOf p -> PredSet p -> Bool
  notMember k pm = IntSet.notMember (one k) (predSet pm)
  
  union :: PredSet p -> PredSet p -> PredSet p
  union x y = PredSet $ IntSet.union (predSet x) (predSet y)
  
  difference :: PredSet p -> PredSet p -> PredSet p
  difference x y = PredSet $ IntSet.difference (predSet x) (predSet y)
  
  -- -----------------------------------------------------------------------------
  
  keysSet :: PredMap p v -> PredSet p
  keysSet = PredSet . IntMap.keysSet . PredMap.predMap
  
  fromSet :: (IdOf p -> v) -> PredSet p -> PredMap p v
  fromSet f =
    PredMap.PredMap . IntMap.fromSet (f . IdOf . Fid . fromIntegral) . predSet
  
  fromSetM :: Monad m => (IdOf p -> m v) -> PredSet p -> m (PredMap p v)
  fromSetM f ps = do
    let ks = toAscList ps
    kvs <- mapM (\ k -> (k,) <$> f k) ks
    return (PredMap.fromDistinctAscList kvs)
  