-- | Small sets represented as immutable bit arrays for fast membership checking.
--
-- Membership checking is O(1), but all other operations are O(n)
-- where n is the size of the element type.
-- The element type needs to implement 'Bounded' and 'Ix'.
--
-- 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()
  , Ix
  , (\\)
  , complement
  , delete
  , difference
  , elems
  , empty
  , fromList, fromAscList, fromDistinctAscList
  , insert
  , intersection
  , isSubsetOf
  , mapMembership
  , member
  , notMember
  , null
  , singleton
  , toList, toAscList
  , total
  , union
  , zipMembershipWith
  ) where

import Prelude hiding (null)

import Control.DeepSeq

import Data.Array.IArray (Ix, Array)
import qualified Data.Array.IArray as Array

-- Note: we might want to use unboxed arrays, but they have no Data instance
--
-- Update: There is currently no need for a Data instance. An attempt
-- was made to replace Array with Data.Array.Unboxed.UArray. Limited
-- testing suggested that this does not make much of a difference in
-- practice (at least not when it comes to type-checking the standard
-- library up to and including Data.Nat, with Agda compiled without
-- -foptimise-heavily).

-- | Let @n@ be the size of type @a@.
type SmallSetElement a = (Bounded a, Ix a)
newtype SmallSet a = SmallSet { forall a. SmallSet a -> Array a Bool
theSmallSet :: Array a Bool }
  deriving (SmallSet a -> SmallSet a -> Bool
(SmallSet a -> SmallSet a -> Bool)
-> (SmallSet a -> SmallSet a -> Bool) -> Eq (SmallSet a)
forall a. Ix a => SmallSet a -> SmallSet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
== :: SmallSet a -> SmallSet a -> Bool
$c/= :: forall a. Ix 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 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. Ix a => Eq (SmallSet a)
forall a. Ix a => SmallSet a -> SmallSet a -> Bool
forall a. Ix a => SmallSet a -> SmallSet a -> Ordering
forall a. Ix a => SmallSet a -> SmallSet a -> SmallSet a
$ccompare :: forall a. Ix a => SmallSet a -> SmallSet a -> Ordering
compare :: SmallSet a -> SmallSet a -> Ordering
$c< :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
< :: SmallSet a -> SmallSet a -> Bool
$c<= :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
<= :: SmallSet a -> SmallSet a -> Bool
$c> :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
> :: SmallSet a -> SmallSet a -> Bool
$c>= :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
>= :: SmallSet a -> SmallSet a -> Bool
$cmax :: forall a. Ix a => SmallSet a -> SmallSet a -> SmallSet a
max :: SmallSet a -> SmallSet a -> SmallSet a
$cmin :: forall a. Ix 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. (Ix a, Show a) => Int -> SmallSet a -> ShowS
forall a. (Ix a, Show a) => [SmallSet a] -> ShowS
forall a. (Ix a, Show a) => SmallSet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. (Ix a, Show a) => Int -> SmallSet a -> ShowS
showsPrec :: Int -> SmallSet a -> ShowS
$cshow :: forall a. (Ix a, Show a) => SmallSet a -> String
show :: SmallSet a -> String
$cshowList :: forall a. (Ix a, Show a) => [SmallSet a] -> ShowS
showList :: [SmallSet a] -> ShowS
Show, SmallSet a -> ()
(SmallSet a -> ()) -> NFData (SmallSet a)
forall a. NFData a => SmallSet a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => SmallSet a -> ()
rnf :: SmallSet a -> ()
NFData)

-- * Query

-- | Time O(n)!
null :: SmallSetElement a => SmallSet a -> Bool
null :: forall a. SmallSetElement a => SmallSet a -> Bool
null = (Bool -> Bool) -> [Bool] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
False) ([Bool] -> Bool) -> (SmallSet a -> [Bool]) -> SmallSet a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a Bool -> [Bool]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
Array.elems (Array a Bool -> [Bool])
-> (SmallSet a -> Array a Bool) -> SmallSet a -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallSet a -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet

-- | 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 -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet SmallSet a
s Array a Bool -> a -> Bool
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
Array.! 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

-- | Time O(n).
isSubsetOf ::  SmallSetElement a => SmallSet a -> SmallSet a -> Bool
isSubsetOf :: forall a. SmallSetElement a => SmallSet a -> SmallSet a -> Bool
isSubsetOf SmallSet a
s SmallSet a
t = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
toBoolListZipWith Bool -> Bool -> Bool
implies SmallSet a
s SmallSet a
t
  where implies :: Bool -> Bool -> Bool
implies Bool
a Bool
b = if Bool
a then Bool
b else Bool
True

-- * Construction

-- | The empty set.  Time O(n).
empty :: SmallSetElement a => SmallSet a
empty :: forall a. SmallSetElement a => SmallSet a
empty = [Bool] -> SmallSet a
forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList (Bool -> [Bool]
forall a. a -> [a]
repeat Bool
False)

-- | The full set.  Time O(n).
total :: SmallSetElement a => SmallSet a
total :: forall a. SmallSetElement a => SmallSet a
total = [Bool] -> SmallSet a
forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList (Bool -> [Bool]
forall a. a -> [a]
repeat Bool
True)

-- | A singleton set.  Time O(n).
singleton :: SmallSetElement a => a -> SmallSet a
singleton :: forall a. SmallSetElement a => a -> SmallSet a
singleton a
a = [a] -> SmallSet a
forall a. SmallSetElement a => [a] -> SmallSet a
fromList [a
a]

-- | Time O(n).
insert :: SmallSetElement a => a -> SmallSet a -> SmallSet a
insert :: forall a. SmallSetElement a => a -> SmallSet a -> SmallSet a
insert a
a = [(a, Bool)] -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
[(a, Bool)] -> SmallSet a -> SmallSet a
update [(a
a,Bool
True)]

-- | Time O(n).
delete :: SmallSetElement a => a -> SmallSet a -> SmallSet a
delete :: forall a. SmallSetElement a => a -> SmallSet a -> SmallSet a
delete a
a = [(a, Bool)] -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
[(a, Bool)] -> SmallSet a -> SmallSet a
update [(a
a,Bool
False)]

-- * Combine

-- | Time O(n).
complement :: SmallSetElement a => SmallSet a -> SmallSet a
complement :: forall a. SmallSetElement a => SmallSet a -> SmallSet a
complement = (Bool -> Bool) -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
(Bool -> Bool) -> SmallSet a -> SmallSet a
mapMembership Bool -> Bool
not

-- | Time O(n).
difference, (\\) :: SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
difference :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
difference = (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith ((Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a)
-> (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
forall a b. (a -> b) -> a -> b
$ \ Bool
b Bool
c -> Bool
b Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
c
\\ :: 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(n).
intersection ::  SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
intersection :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
intersection = (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith Bool -> Bool -> Bool
(&&)

-- | Time O(n).
union ::  SmallSetElement a => SmallSet a -> SmallSet a -> SmallSet a
union :: forall a.
SmallSetElement a =>
SmallSet a -> SmallSet a -> SmallSet a
union = (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith Bool -> Bool -> Bool
(||)

-- | Time O(n).
mapMembership :: SmallSetElement a => (Bool -> Bool) -> SmallSet a -> SmallSet a
mapMembership :: forall a.
SmallSetElement a =>
(Bool -> Bool) -> SmallSet a -> SmallSet a
mapMembership Bool -> Bool
f = Array a Bool -> SmallSet a
forall a. Array a Bool -> SmallSet a
SmallSet (Array a Bool -> SmallSet a)
-> (SmallSet a -> Array a Bool) -> SmallSet a -> SmallSet a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Bool) -> Array a Bool -> Array a Bool
forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
Array.amap Bool -> Bool
f (Array a Bool -> Array a Bool)
-> (SmallSet a -> Array a Bool) -> SmallSet a -> Array a Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallSet a -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet

-- | Time O(n).
zipMembershipWith :: SmallSetElement a => (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith :: forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith Bool -> Bool -> Bool
f SmallSet a
s SmallSet a
t = [Bool] -> SmallSet a
forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList ([Bool] -> SmallSet a) -> [Bool] -> SmallSet a
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
toBoolListZipWith Bool -> Bool -> Bool
f SmallSet a
s SmallSet a
t

-- * Conversion

-- | Time O(n).
elems, toList, toAscList :: SmallSetElement a => SmallSet a -> [a]
elems :: forall a. SmallSetElement a => SmallSet a -> [a]
elems     = ((a, Bool) -> a) -> [(a, Bool)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, Bool) -> a
forall a b. (a, b) -> a
fst ([(a, Bool)] -> [a])
-> (SmallSet a -> [(a, Bool)]) -> SmallSet a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, Bool) -> Bool) -> [(a, Bool)] -> [(a, Bool)]
forall a. (a -> Bool) -> [a] -> [a]
filter (a, Bool) -> Bool
forall a b. (a, b) -> b
snd ([(a, Bool)] -> [(a, Bool)])
-> (SmallSet a -> [(a, Bool)]) -> SmallSet a -> [(a, Bool)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a Bool -> [(a, Bool)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
Array.assocs (Array a Bool -> [(a, Bool)])
-> (SmallSet a -> Array a Bool) -> SmallSet a -> [(a, Bool)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallSet a -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet
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            = ([(a, Bool)] -> SmallSet a -> SmallSet a)
-> SmallSet a -> [(a, Bool)] -> SmallSet a
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(a, Bool)] -> SmallSet a -> SmallSet a
forall a.
SmallSetElement a =>
[(a, Bool)] -> SmallSet a -> SmallSet a
update SmallSet a
forall a. SmallSetElement a => SmallSet a
empty ([(a, Bool)] -> SmallSet a)
-> ([a] -> [(a, Bool)]) -> [a] -> SmallSet a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (a, Bool)) -> [a] -> [(a, Bool)]
forall a b. (a -> b) -> [a] -> [b]
map (,Bool
True)
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

-- * Internal

-- | Time O(n).  Assumes @Bool@-vector of length @n@.
fromBoolList :: SmallSetElement a => [Bool] -> SmallSet a
fromBoolList :: forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList = Array a Bool -> SmallSet a
forall a. Array a Bool -> SmallSet a
SmallSet (Array a Bool -> SmallSet a)
-> ([Bool] -> Array a Bool) -> [Bool] -> SmallSet a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a) -> [Bool] -> Array a Bool
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
Array.listArray (a
forall a. Bounded a => a
minBound, a
forall a. Bounded a => a
maxBound)

-- | Time O(n).  Produces @Bool@-vector of length @n@.
toBoolList :: SmallSetElement a => SmallSet a -> [Bool]
toBoolList :: forall a. SmallSetElement a => SmallSet a -> [Bool]
toBoolList = Array a Bool -> [Bool]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
Array.elems (Array a Bool -> [Bool])
-> (SmallSet a -> Array a Bool) -> SmallSet a -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallSet a -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet

-- | Time O(n).  Produces @Bool@-vector of length @n@.
toBoolListZipWith :: SmallSetElement a => (Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
toBoolListZipWith :: forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> [Bool]
toBoolListZipWith Bool -> Bool -> Bool
f SmallSet a
s SmallSet a
t = (Bool -> Bool -> Bool) -> [Bool] -> [Bool] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bool -> Bool -> Bool
f (SmallSet a -> [Bool]
forall a. SmallSetElement a => SmallSet a -> [Bool]
toBoolList SmallSet a
s) (SmallSet a -> [Bool]
forall a. SmallSetElement a => SmallSet a -> [Bool]
toBoolList SmallSet a
t)

-- | Time O(n).  Bulk insert/delete.
update :: SmallSetElement a => [(a,Bool)] -> SmallSet a -> SmallSet a
update :: forall a.
SmallSetElement a =>
[(a, Bool)] -> SmallSet a -> SmallSet a
update [(a, Bool)]
u SmallSet a
s = Array a Bool -> SmallSet a
forall a. Array a Bool -> SmallSet a
SmallSet (Array a Bool -> SmallSet a) -> Array a Bool -> SmallSet a
forall a b. (a -> b) -> a -> b
$ SmallSet a -> Array a Bool
forall a. SmallSet a -> Array a Bool
theSmallSet SmallSet a
s Array a Bool -> [(a, Bool)] -> Array a Bool
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
Array.// [(a, Bool)]
u