-- | 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
forall a. Ix a => SmallSet a -> SmallSet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: 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
Eq, 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
min :: SmallSet a -> SmallSet a -> SmallSet a
$cmin :: forall a. Ix a => SmallSet a -> SmallSet a -> SmallSet a
max :: SmallSet a -> SmallSet a -> SmallSet a
$cmax :: forall a. Ix a => SmallSet a -> SmallSet a -> SmallSet a
>= :: 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
$c< :: forall a. Ix a => SmallSet a -> SmallSet a -> Bool
compare :: SmallSet a -> SmallSet a -> Ordering
$ccompare :: forall a. Ix a => SmallSet a -> SmallSet a -> Ordering
Ord, Int -> SmallSet a -> ShowS
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
showList :: [SmallSet a] -> ShowS
$cshowList :: forall a. (Ix a, Show a) => [SmallSet a] -> ShowS
show :: SmallSet a -> String
$cshow :: forall a. (Ix a, Show a) => SmallSet a -> String
showsPrec :: Int -> SmallSet a -> ShowS
$cshowsPrec :: forall a. (Ix a, Show a) => Int -> SmallSet a -> ShowS
Show, SmallSet a -> ()
forall a. NFData a => SmallSet a -> ()
forall a. (a -> ()) -> NFData a
rnf :: SmallSet a -> ()
$crnf :: forall a. NFData a => SmallSet a -> ()
NFData)

-- * Query

-- | Time O(n)!
null :: SmallSetElement a => SmallSet a -> Bool
null :: forall a. SmallSetElement a => SmallSet a -> Bool
null = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (forall a. Eq a => a -> a -> Bool
== Bool
False) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
Array.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall a. SmallSet a -> Array a Bool
theSmallSet SmallSet a
s 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *). Foldable t => t Bool -> Bool
and forall a b. (a -> b) -> a -> b
$ 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 = forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList (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 = forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList (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 = 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 = 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 = 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 = 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 = forall a.
SmallSetElement a =>
(Bool -> Bool -> Bool) -> SmallSet a -> SmallSet a -> SmallSet a
zipMembershipWith 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
(\\)       = 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 = 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 = 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 = forall a. Array a Bool -> SmallSet a
SmallSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall a. SmallSetElement a => [Bool] -> SmallSet a
fromBoolList forall a b. (a -> b) -> a -> b
$ 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     = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
Array.assocs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. SmallSet a -> Array a Bool
theSmallSet
toList :: forall a. SmallSetElement a => SmallSet a -> [a]
toList    = forall a. SmallSetElement a => SmallSet a -> [a]
elems
toAscList :: forall a. SmallSetElement a => SmallSet a -> [a]
toAscList = 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            = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a.
SmallSetElement a =>
[(a, Bool)] -> SmallSet a -> SmallSet a
update forall a. SmallSetElement a => SmallSet a
empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (,Bool
True)
fromAscList :: forall a. SmallSetElement a => [a] -> SmallSet a
fromAscList         = forall a. SmallSetElement a => [a] -> SmallSet a
fromList
fromDistinctAscList :: forall a. SmallSetElement a => [a] -> SmallSet a
fromDistinctAscList = 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 = forall a. Array a Bool -> SmallSet a
SmallSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
Array.listArray (forall a. Bounded a => a
minBound, 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 = forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
Array.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bool -> Bool -> Bool
f (forall a. SmallSetElement a => SmallSet a -> [Bool]
toBoolList SmallSet a
s) (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 = forall a. Array a Bool -> SmallSet a
SmallSet forall a b. (a -> b) -> a -> b
$ forall a. SmallSet a -> Array a Bool
theSmallSet SmallSet a
s forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
Array.// [(a, Bool)]
u