{-# LANGUAGE CPP #-}

{-# OPTIONS_GHC -Wunused-imports #-}

-- | Small sets represented as a bitmask for fast membership checking.
--
-- With the exception of converting to/from lists, all operations are O(1).
--
-- Mimics the interface of 'Data.Set'.
--
-- Import as:
-- @
--    import qualified Agda.Utils.SmallSet as SmallSet
--    import Agda.Utils.SmallSet (SmallSet)
-- @

module Agda.Utils.SmallSet
  ( SmallSet(..)
  , SmallSetElement
  , Ix
  , (\\)
  , complement
  , delete
  , difference
  , elems
  , empty
  , fromList, fromAscList, fromDistinctAscList
  , insert
  , intersection
  , member
  , notMember
  , null
  , singleton
  , toList, toAscList
  , total
  , union
  ) where

import Prelude hiding (null)

import Control.DeepSeq

import Data.Word (Word64)
#if !MIN_VERSION_base(4,20,0)
import Data.List (foldl')
#endif
import Data.Bits hiding (complement)
import qualified Data.Bits as Bits
import Data.Ix

import qualified Agda.Utils.Null as Null

-- | An element in a small set.
--
-- This must implement 'Bounded' and 'Ix', and contain at most 64 values.
class (Bounded a, Ix a) => SmallSetElement a where

newtype SmallSet a = SmallSet { forall a. SmallSet a -> Word64
theSmallSet :: Word64 }
  deriving (SmallSet a -> SmallSet a -> Bool
(SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> Bool) -> Eq (SmallSet a)
forall a. SmallSet a -> SmallSet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. SmallSet a -> SmallSet a -> Bool
== :: SmallSet a -> SmallSet a -> Bool
$c/= :: forall a. SmallSet a -> SmallSet a -> Bool
/= :: SmallSet a -> SmallSet a -> Bool
Eq, Eq (SmallSet a)
Eq (SmallSet a) =>
(SmallSet a -> SmallSet a -> Ordering)
-> (SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> SmallSet a)
-> (SmallSet a -> SmallSet a -> SmallSet a)
-> Ord (SmallSet a)
SmallSet a -> SmallSet a -> Bool
SmallSet a -> SmallSet a -> Ordering
SmallSet a -> SmallSet a -> SmallSet a
forall a. Eq (SmallSet a)
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. SmallSet a -> SmallSet a -> Bool
forall a. SmallSet a -> SmallSet a -> Ordering
forall a. SmallSet a -> SmallSet a -> SmallSet a
$ccompare :: forall a. SmallSet a -> SmallSet a -> Ordering
compare :: SmallSet a -> SmallSet a -> Ordering
$c< :: forall a. SmallSet a -> SmallSet a -> Bool
< :: SmallSet a -> SmallSet a -> Bool
$c<= :: forall a. SmallSet a -> SmallSet a -> Bool
<= :: SmallSet a -> SmallSet a -> Bool
$c> :: forall a. SmallSet a -> SmallSet a -> Bool
> :: SmallSet a -> SmallSet a -> Bool
$c>= :: forall a. SmallSet a -> SmallSet a -> Bool
>= :: SmallSet a -> SmallSet a -> Bool
$cmax :: forall a. SmallSet a -> SmallSet a -> SmallSet a
max :: SmallSet a -> SmallSet a -> SmallSet a
$cmin :: forall a. SmallSet a -> SmallSet a -> SmallSet a
min :: SmallSet a -> SmallSet a -> SmallSet a
Ord, Int -> SmallSet a -> ShowS
[SmallSet a] -> ShowS
SmallSet a -> String
(Int -> SmallSet a -> ShowS)
-> (SmallSet a -> String)
-> ([SmallSet a] -> ShowS)
-> Show (SmallSet a)
forall a. Int -> SmallSet a -> ShowS
forall a. [SmallSet a] -> ShowS
forall a. SmallSet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Int -> SmallSet a -> ShowS
showsPrec :: Int -> SmallSet a -> ShowS
$cshow :: forall a. SmallSet a -> String
show :: SmallSet a -> String
$cshowList :: forall a. [SmallSet a] -> ShowS
showList :: [SmallSet a] -> ShowS
Show, SmallSet a -> ()
(SmallSet a -> ()) -> NFData (SmallSet a)
forall a. SmallSet a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. SmallSet a -> ()
rnf :: SmallSet a -> ()
NFData)


instance SmallSetElement a => Null.Null (SmallSet a) where
  empty :: SmallSet a
empty = SmallSet a
forall a. SmallSetElement a => SmallSet a
empty
  null :: SmallSet a -> Bool
null = SmallSet a -> Bool
forall a. SmallSetElement a => SmallSet a -> Bool
null

instance SmallSetElement a => Semigroup (SmallSet a) where
  <> :: SmallSet a -> SmallSet a -> SmallSet a
(<>) = SmallSet a -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
union

instance SmallSetElement a => Monoid (SmallSet a) where
  mempty :: SmallSet a
mempty = SmallSet a
forall a. SmallSetElement a => SmallSet a
empty

-- * Query

-- | Time O(1).
null :: SmallSetElement a => SmallSet a -> Bool
null :: forall a. SmallSetElement a => SmallSet a -> Bool
null SmallSet a
s = SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0

-- | Time O(1).
member :: SmallSetElement a => a -> SmallSet a -> Bool
member :: forall a. SmallSetElement a => a -> SmallSet a -> Bool
member a
a SmallSet a
s = SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` a -> Int
forall a. SmallSetElement a => a -> Int
idx a
a

-- | @not . member a@.  Time O(1).
notMember :: SmallSetElement a => a -> SmallSet a -> Bool
notMember :: forall a. SmallSetElement a => a -> SmallSet a -> Bool
notMember a
a = Bool -> Bool
not (Bool -> Bool) -> (SmallSet a -> Bool) -> SmallSet a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> SmallSet a -> Bool
forall a. SmallSetElement a => a -> SmallSet a -> Bool
member a
a

-- * Construction

-- | The empty set.  Time O(1).
empty :: SmallSetElement a => SmallSet a
empty :: forall a. SmallSetElement a => SmallSet a
empty = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet Word64
0

-- | The full set.  Time O(1).
total :: forall a. SmallSetElement a => SmallSet a
total :: forall a. SmallSetElement a => SmallSet a
total = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ Word64 -> Word64
forall a. Bits a => a -> a
Bits.complement Word64
0

-- | A singleton set.  Time O(1).
singleton :: SmallSetElement a => a -> SmallSet a
singleton :: forall a. SmallSetElement a => a -> SmallSet a
singleton a
a = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ Int -> Word64
forall a. Bits a => Int -> a
bit (a -> Int
forall a. SmallSetElement a => a -> Int
idx a
a)

-- | Time O(1).
insert :: SmallSetElement a => a -> SmallSet a -> SmallSet a
insert :: forall a. SmallSetElement a => a -> SmallSet a -> SmallSet a
insert a
a SmallSet a
s = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`setBit` a -> Int
forall a. SmallSetElement a => a -> Int
idx a
a

-- | Time O(1).
delete :: SmallSetElement a => a -> SmallSet a -> SmallSet a
delete :: forall a. SmallSetElement a => a -> SmallSet a -> SmallSet a
delete a
a SmallSet a
s = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`clearBit` a -> Int
forall a. SmallSetElement a => a -> Int
idx a
a

-- * Combine

-- | Time O(n).
complement :: SmallSetElement a => SmallSet a -> SmallSet a
complement :: forall a. SmallSetElement a => SmallSet a -> SmallSet a
complement = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a)
-> (SmallSet a -> Word64) -> SmallSet a -> SmallSet a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Word64
forall a. Bits a => a -> a
Bits.complement (Word64 -> Word64)
-> (SmallSet a -> Word64) -> SmallSet a -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet

-- | Time O(1).
difference, (\\) :: SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
difference :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
difference SmallSet a
s SmallSet a
t = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64 -> Word64
forall a. Bits a => a -> a
Bits.complement (SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
t)
\\ :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
(\\)       = SmallSet a -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
difference

-- | Time O(1).
intersection ::  SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
intersection :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
intersection SmallSet a
s SmallSet a
t = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
t

-- | Time O(n).
union ::  SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
union :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
union SmallSet a
s SmallSet a
t = Word64 -> SmallSet a
forall a. Word64 -> SmallSet a
SmallSet (Word64 -> SmallSet a) -> Word64 -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
t

-- * Conversion

-- | Time O(n).
elems, toList, toAscList :: SmallSetElement a => SmallSet a -> [a]
elems :: forall a. SmallSetElement a => SmallSet a -> [a]
elems SmallSet a
s   = (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter (\a
i -> SmallSet a -> Word64
forall a. SmallSet a -> Word64
theSmallSet SmallSet a
s Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` a -> Int
forall a. SmallSetElement a => a -> Int
idx a
i) ((a, a) -> [a]
forall a. Ix a => (a, a) -> [a]
range (a, a)
forall a. SmallSetElement a => (a, a)
bounds)
toList :: forall a. SmallSetElement a => SmallSet a -> [a]
toList    = SmallSet a -> [a]
forall a. SmallSetElement a => SmallSet a -> [a]
elems
toAscList :: forall a. SmallSetElement a => SmallSet a -> [a]
toAscList = SmallSet a -> [a]
forall a. SmallSetElement a => SmallSet a -> [a]
elems

-- | Time O(n).
fromList, fromAscList, fromDistinctAscList :: SmallSetElement a => [a] -> SmallSet a
fromList :: forall a. SmallSetElement a => [a] -> SmallSet a
fromList            = (SmallSet a -> a -> SmallSet a) -> SmallSet a -> [a] -> SmallSet a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((a -> SmallSet a -> SmallSet a) -> SmallSet a -> a -> SmallSet a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> SmallSet a -> SmallSet a
forall a. SmallSetElement a => a -> SmallSet a -> SmallSet a
insert) SmallSet a
forall a. SmallSetElement a => SmallSet a
empty
fromAscList :: forall a. SmallSetElement a => [a] -> SmallSet a
fromAscList         = [a] -> SmallSet a
forall a. SmallSetElement a => [a] -> SmallSet a
fromList
fromDistinctAscList :: forall a. SmallSetElement a => [a] -> SmallSet a
fromDistinctAscList = [a] -> SmallSet a
forall a. SmallSetElement a => [a] -> SmallSet a
fromList

-- * Internals

bounds :: SmallSetElement a => (a, a)
bounds :: forall a. SmallSetElement a => (a, a)
bounds = (a
forall a. Bounded a => a
minBound, a
forall a. Bounded a => a
maxBound)

idx :: SmallSetElement a => a -> Int
idx :: forall a. SmallSetElement a => a -> Int
idx a
a = (a, a) -> a -> Int
forall a. Ix a => (a, a) -> a -> Int
index (a, a)
forall a. SmallSetElement a => (a, a)
bounds a
a