-- Needed to get haddock to pick up on identifiers with a #.
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE CPP #-}

{-|
Manage sets of natural numbers (de Bruijn indices).

This file contains an implementation of finite sets of integers
that is optimized for storing sets of free variables. Notably,
de Bruijn indices/levels are not uniformly distributed across the
range of 'Int', and are always very small (EG: less than 64).

This makes 'Data.IntSet.IntSet' somewhat ill-suited for the task.
'Data.IntSet.IntSet' is designed to be able to handle values distributed
across the entire range of 'Int', at the cost of pointer indirections
and poor memory locality.

Instead, 'VarSet' stores free variables as a bitvector. When the
elements in the set are all smaller than 64, we can fit the entire set
into a single unboxed 'GHC.Base.Word#', which makes most of our set operations
a single instruction.

When we exceed this bound, we fall back to an unboxed 'GHC.Base.ByteArray#'.
Experimentally, this happens less than 0.0001% of the time.

= Experimental Results

Agda builds instrumented with 'traceVarSetSize' give the following
statistics for 'VarSet' sizes across the following libraries:

+-------------+------------------------------------------+-------------------+---------------------+
| library     | commit hash                              | varsets allocated | non-word sized sets |
+=============+==========================================+===================+=====================+
| 1lab        | 11f4672ca7dcfdccb1b4ce94ca4ca42c9b732e62 | 118198551         | 0                   |
+-------------+------------------------------------------+-------------------+---------------------+
| agda-stdlib | 0e97e2ed1f999ea0d2cce2a3b2395d5a9a7bd36a | 38960049          | 2178                |
+-------------+------------------------------------------+-------------------+---------------------+
| cubical     | 1f2af52701945bef003a525f78fa41aeadb7c6ae | 82751805          | 0                   |
+-------------+------------------------------------------+-------------------+---------------------+
-}

module Agda.Utils.VarSet
  (
  -- * VarSet type
  --
  -- $varSetType
    VarSet(..)
  -- * Construction
  , empty
  , singleton
  , fromList
  , fromDescList
  , full
  , range
  -- * Insertion
  , insert
  , inserts
  , insertsDesc
  -- * Deletion
  , delete
  -- * Queries
  , null
  , member
  , lookupMin
  , lookupMax
  , size
  , disjoint
  , isSubsetOf
  , inRange
  -- * Combining variable sets
  , union
  , unions
  , intersection
  , difference
  , (\\)
  , complement
  -- * Filters
  , split
  , filterLT
  , filterLE
  , filterGT
  , filterGE
  -- * Folds
  --
  -- $varSetFolds
  , foldr
  , foldl
  , foldr'
  , foldl'
  -- * Views
  , minView
  , maxView
  -- * Contextual operations
  , strengthen
  , weaken
  -- * Conversions
  , toDescList
  , toAscList
  )
  where

-- We need the machines word size for some bitwise operations.
#include "MachDeps.h"

import Control.DeepSeq
import Control.Monad

import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as BS
import qualified Data.Foldable as Fold
import Data.Hashable

import GHC.Base hiding (empty, foldr)
import GHC.Num.BigNat
import GHC.Num.WordArray

import Debug.Trace

import Prelude (Num(..), Show(..))
import qualified Prelude as P

import Agda.Utils.ByteArray
import Agda.Utils.Word

-- $varSetType
-- = Invariants
-- A variable set whose maximum entry is less than than 64 is always encoded
-- via the 'VS#' constructor.
-- The final word of a 'VB#' is always non-zero.
--
-- = Performance
-- GHC will use pointer-tagging for datatypes with less than 4 constructors
-- on 32 bit machines (8 on 64 bit), which means that 'VarSet' will be laid
-- out as a tagged pointer to a word-sized piece of data.
--
-- We also pay attention to the constructor ordering here.
-- GHC will lay out the machine code for case statements in
-- reverse order of the datatypes constructors. Both ARM and x86 statically
-- predict that forward branches are not taken, so putting 'VB#' before
-- 'VS#' means that we can eek out a few extra drops of performance.
--
-- Finally, many functions in this module follow the following pattern:
-- @
-- union :: VarSet -> VarSet -> VarSet
-- union (VS# w1) (VS# w2) = VS# (w1 `or#` w2)
-- union vs1 vs2 = unionSlow vs1 vs2
-- {-# INLINE union #-}
--
-- unionSlow :: VarSet -> VarSet -> VarSet
-- unionSlow (VS# w1) (VS# w2) = VS# (w1 `or#` w2)
-- unionSlow (VS# w1) (VB# bs2) = VB# (bigNatOrWord# bs2 w1)
-- unionSlow (VB# bs1) (VS# w2) = VB# (bigNatOrWord# bs1 w2)
-- unionSlow (VB# bs1) (VB# bs2) = VB# (bigNatOr bs1 bs2)
-- {-# NOINLINE unionSlow #-}
-- @
--
-- This makes the 0.0001% of slow cases notably worse, but makes the
-- happy path approximately 3x faster for things like 'inRange', as
-- GHC is able to coalesce all of the happy paths.

-- | A set of de Bruijn indices/levels.
data VarSet
  = VB# ByteArray#
  -- ^ A variable set whose maximum entry is greater than or equal to 64.
  | VS# Word#
  -- ^ A small variable set whose maximum entry is less than than 64.

-- | Convert a 'ByteArray#' to a variable set, upholding
-- the invariant that all word-sized variable sets should use
-- 'VS#'.
byteArrayToVarSet# :: ByteArray# -> VarSet
byteArrayToVarSet# :: ByteArray# -> VarSet
byteArrayToVarSet# ByteArray#
bs =
  case ByteArray# -> Int#
bigNatSize# ByteArray#
bs of
    Int#
0# -> VarSet
empty -- Dont allocate if we don't have to.
    Int#
1# -> Word# -> VarSet
VS# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs Int#
0#)
    Int#
_  -> ByteArray# -> VarSet
VB# ByteArray#
bs

instance Eq VarSet where
  (VS# Word#
xs) == :: VarSet -> VarSet -> Bool
== (VS# Word#
ys) = Int# -> Bool
isTrue# (Word#
xs Word# -> Word# -> Int#
`eqWord#` Word#
ys)
  (VB# ByteArray#
_) == (VS# Word#
_) = Bool
False
  (VS# Word#
_) == (VB# ByteArray#
_) = Bool
False
  (VB# ByteArray#
xs) == (VB# ByteArray#
ys) = ByteArray# -> ByteArray# -> Bool
bigNatEq ByteArray#
xs ByteArray#
ys

instance Ord VarSet where
  compare :: VarSet -> VarSet -> Ordering
compare (VS# Word#
xs) (VS# Word#
ys) = Word# -> Word# -> Ordering
compareWord# Word#
xs Word#
ys
  compare (VB# ByteArray#
_) (VS# Word#
_) = Ordering
GT
  compare (VS# Word#
_) (VB# ByteArray#
_) = Ordering
LT
  compare (VB# ByteArray#
xs) (VB# ByteArray#
ys) = ByteArray# -> ByteArray# -> Ordering
bigNatCompare ByteArray#
xs ByteArray#
ys

instance P.Show VarSet where
  -- Use the same direction as the 'Show' instance for 'IntSet'.
  showsPrec :: Int -> VarSet -> ShowS
showsPrec Int
d VarSet
vs =
    Bool -> ShowS -> ShowS
P.showParen (Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
    String -> ShowS
P.showString String
"fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> ShowS
forall a. Show a => a -> ShowS
P.shows (VarSet -> [Int]
toAscList VarSet
vs)

instance NFData VarSet where
  rnf :: VarSet -> ()
rnf (VS# Word#
_) = ()
  rnf (VB# ByteArray#
_) = ()

-- | @(<>)@ is @'union'@
instance Semigroup VarSet where
  <> :: VarSet -> VarSet -> VarSet
(<>) = VarSet -> VarSet -> VarSet
union

-- | @mempty@ is @'empty'@.
instance Monoid VarSet where
  mempty :: VarSet
mempty = VarSet
empty

-- This instance is a bit suboptimal.
-- Ideally, 'hashWithSalt' should directly just call an unboxed 'mixHash' when we have a 'VS#',
-- but 'Hashable' doesn't expose enough of its internals for this.
instance Hashable VarSet where
  hash :: VarSet -> Int
hash (VS# Word#
w) = Int# -> Int
I# (Word# -> Int#
word2Int# Word#
w)
  hash (VB# ByteArray#
bs) = ByteArray# -> Int -> Int -> Int
hashByteArray ByteArray#
bs Int
0 (Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
bs))

  hashWithSalt :: Int -> VarSet -> Int
hashWithSalt Int
salt (VS# Word#
w)  = Int -> Int -> Int
forall a. Hashable a => Int -> a -> Int
hashWithSalt (Int# -> Int
I# (Word# -> Int#
word2Int# Word#
w)) Int
salt
  hashWithSalt Int
salt (VB# ByteArray#
bs) = ByteArray# -> Int -> Int -> Int -> Int
hashByteArrayWithSalt ByteArray#
bs Int
0 (Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
bs)) Int
salt

-- Could be much more efficient. Ideally, we would just have our hands
-- on a pointer and an offset, but 'Binary' doesn't give a nice interface
-- for this.
instance Binary VarSet where
  put :: VarSet -> Put
put (VS# Word#
w) = Word8 -> Put
putWord8 Word8
0 Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Word -> Put
forall t. Binary t => t -> Put
put (Word# -> Word
W# Word#
w)
  put (VB# ByteArray#
bs) = Word8 -> Put
putWord8 Word8
1 Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Int -> Put
forall t. Binary t => t -> Put
put (Int# -> Int
I# Int#
len) Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Int# -> Put
loop (Int#
len Int# -> Int# -> Int#
-# Int#
1#)
    where
      len :: Int#
len = ByteArray# -> Int#
bigNatSize# ByteArray#
bs

      loop :: Int# -> Put
      loop :: Int# -> Put
loop Int#
i =
        if Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
>=# Int#
0#) then
          Word -> Put
forall t. Binary t => t -> Put
put (Word# -> Word
W# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs Int#
i)) Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Int# -> Put
loop (Int#
i Int# -> Int# -> Int#
-# Int#
1#)
        else
          Put
forall a. Monoid a => a
mempty

  get :: Get VarSet
get = do
    tag <- forall t. Binary t => Get t
get @Word8
    case tag of
      Word8
0 -> do
        W# w <- forall t. Binary t => Get t
get @Word
        pure (VS# w)
      Word8
_ -> do
        -- This is quite bad, but we don't have very good tools for
        -- doing this efficiently. This should really never happen
        -- though, so it's not the end of the world.
        len <- forall t. Binary t => Get t
get @Int
        words <- replicateM len (get @Word)
        pure (VB# (bigNatFromWordList words))

--------------------------------------------------------------------------------
-- Construction

-- | The empty variable set.
empty :: VarSet
empty :: VarSet
empty = Word# -> VarSet
VS# Word#
0##
{-# INLINE CONLIKE empty #-}

-- | A variable set with only a single entry.
singleton :: Int -> VarSet
singleton :: Int -> VarSet
singleton (I# Int#
i) =
  if Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) then
    Word# -> VarSet
VS# (Int# -> Word#
uncheckedBitWord# Int#
i)
  else
    ByteArray# -> VarSet
VB# (Word# -> ByteArray#
bigNatBit# (Int# -> Word#
int2Word# Int#
i))
{-# INLINABLE singleton #-}

-- | Construct a set of variables from a list of indices.
fromList :: [Int] -> VarSet
fromList :: [Int] -> VarSet
fromList [Int]
is = [Int] -> VarSet -> VarSet
inserts [Int]
is VarSet
empty
{-# INLINE fromList #-}

-- | Construct a set of variables from a list of indices sorted in descending order.
-- This can be more efficient than 'fromList', as we only need to perform a single
-- to determine that elements of the list are smaller than 64.
--
-- The precondition that the variables are sorted in descending order is
-- not checked by 'fromDescList'.
fromDescList :: [Int] -> VarSet
fromDescList :: [Int] -> VarSet
fromDescList [Int]
is = [Int] -> VarSet -> VarSet
insertsDesc [Int]
is VarSet
empty

-- | Construct a variable set that contains the indices @[0..n-1]@.
full :: Int -> VarSet
full :: Int -> VarSet
full (I# Int#
n) =
  if Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
<=# Int#
0#) then
    VarSet
empty
  else if Int# -> Bool
isTrue# (Int#
n Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) then
    Word# -> VarSet
VS# (Int# -> Word#
uncheckedWordOnes# Int#
n)
  else
    ByteArray# -> VarSet
VB# (Int# -> ByteArray#
byteArrayOnes# Int#
n)
{-# INLINABLE full #-}

-- | Construct a variable set that contains the indices @[lo+1..hi-1]@.
range :: Int -> Int -> VarSet
range :: Int -> Int -> VarSet
range Int
lo Int
hi =
  -- Most of the time, @hi@ is the length of the context
  -- and @lo@ is some variable, so it would be a waste of time
  -- to try and avoid work by checking if hi <= lo.
  Int -> VarSet
full Int
hi VarSet -> VarSet -> VarSet
\\ Int -> VarSet
full (Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE range #-}

--------------------------------------------------------------------------------
-- Insertion

-- | Insert an index into a variable set.
insert :: Int -> VarSet -> VarSet
insert :: Int -> VarSet -> VarSet
insert (I# Int#
i) (VS# Word#
w) =
  if Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) then
    Word# -> VarSet
VS# (Word# -> Int# -> Word#
uncheckedSetBitWord# Word#
w Int#
i)
  else
    ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatSetBit# (Word# -> ByteArray#
bigNatFromWord# Word#
w) (Int# -> Word#
int2Word# Int#
i))
insert (I# Int#
i) (VB# ByteArray#
bs) = ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatSetBit# ByteArray#
bs (Int# -> Word#
int2Word# Int#
i))

-- | Insert a list of variables into a @VarSet@.
inserts :: [Int] -> VarSet -> VarSet
inserts :: [Int] -> VarSet -> VarSet
inserts [Int]
xs VarSet
vs = (VarSet -> Int -> VarSet) -> VarSet -> [Int] -> VarSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Fold.foldl' ((Int -> VarSet -> VarSet) -> VarSet -> Int -> VarSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> VarSet -> VarSet
insert) VarSet
vs [Int]
xs
{-# INLINE inserts #-}

-- | Insert a list of variables sorted in descending order into a 'VarSet'.
-- This can be more efficient than 'inserts', as we only need to perform a single
-- to determine that elements of the list are smaller than 64.
--
-- The precondition that the variables are sorted in descending order is
-- not checked by 'insertsDesc'.
insertsDesc :: [Int] -> VarSet -> VarSet
insertsDesc :: [Int] -> VarSet -> VarSet
insertsDesc [] VarSet
vs = VarSet
vs
insertsDesc (I# Int#
x:[Int]
xs) (VS# Word#
w) | Int# -> Bool
isTrue# (Int#
x Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) = VS# (insertsDescFast# xs (uncheckedSetBitWord# w x))
insertsDesc [Int]
xs VarSet
vs = [Int] -> VarSet -> VarSet
inserts [Int]
xs VarSet
vs -- Could be more optimal but this is basically never taken.

-- | Fast path for 'insertsDesc' that omits bounds checking.
insertsDescFast# :: [Int] -> Word# -> Word#
insertsDescFast# :: [Int] -> Word# -> Word#
insertsDescFast# [] Word#
w = Word#
w
insertsDescFast# (I# Int#
x:[Int]
xs) Word#
w = [Int] -> Word# -> Word#
insertsDescFast# [Int]
xs (Word# -> Int# -> Word#
uncheckedSetBitWord# Word#
w Int#
x)

--------------------------------------------------------------------------------
-- Deletion

-- | Delete an index from a variable set.
delete :: Int -> VarSet -> VarSet
delete :: Int -> VarSet -> VarSet
delete (I# Int#
i) vs :: VarSet
vs@(VS# Word#
w) =
  if Int# -> Bool
isTrue# (Int#
i Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) then
    Word# -> VarSet
VS# (Word# -> Int# -> Word#
uncheckedClearBitWord# Word#
w Int#
i)
  else
    VarSet
vs
delete (I# Int#
i) (VB# ByteArray#
vs) = ByteArray# -> VarSet
byteArrayToVarSet# (ByteArray# -> Word# -> ByteArray#
bigNatClearBit# ByteArray#
vs (Int# -> Word#
int2Word# Int#
i))

--------------------------------------------------------------------------------
-- Queries

-- | Is this the empty set?
null :: VarSet -> Bool
null :: VarSet -> Bool
null (VS# Word#
0##) = Bool
True
null VarSet
_ = Bool
False
{-# INLINE null #-}

-- | Is a variable a member of a var set?
member :: Int -> VarSet -> Bool
member :: Int -> VarSet -> Bool
member (I# Int#
i) (VS# Word#
w) = Int# -> Bool
isTrue# ((Int#
0# Int# -> Int# -> Int#
<=# Int#
i) Int# -> Int# -> Int#
`andI#` (Int#
i Int# -> Int# -> Int#
<# WORD_SIZE_IN_BITS#) `andI#` (uncheckedTestBitWord# w i))
member (I# Int#
i) (VB# ByteArray#
bs) = Int# -> Bool
isTrue# (ByteArray# -> Word# -> Int#
bigNatTestBit# ByteArray#
bs (Int# -> Word#
int2Word# Int#
i))
{-# INLINABLE member #-}

-- | Find the smallest index in the variable set.
lookupMin :: VarSet -> Maybe Int
lookupMin :: VarSet -> Maybe Int
lookupMin (VS# Word#
w) =
  if Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) then
    Maybe Int
forall a. Maybe a
Nothing
  else
    Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
lowestBitWord# Word#
w)))
lookupMin (VB# ByteArray#
bs) = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# (Int# -> ByteArray# -> Int#
lookupMinSlow# Int#
0# ByteArray#
bs))

-- | Slow path for 'lookupMin'.
lookupMinSlow# :: Int# -> ByteArray# -> Int#
lookupMinSlow# :: Int# -> ByteArray# -> Int#
lookupMinSlow# Int#
i ByteArray#
bs =
  -- We don't need to check if the index goes out of bounds, as we know that
  -- 0 always gets stored as an @VS#@, so there must be *some* set bit.
  let w :: Word#
w = ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs Int#
i in
  if Int# -> Bool
isTrue# (Word# -> Word# -> Int#
eqWord# Word#
w Word#
0##) then
    Int# -> ByteArray# -> Int#
lookupMinSlow# (Int#
i Int# -> Int# -> Int#
+# Int#
1#) ByteArray#
bs
  else
    WORD_SIZE_IN_BITS# *# i +# (word2Int# (lowestBitWord# w))
{-# NOINLINE lookupMinSlow# #-}

-- | Find the largest index in the variable set.
lookupMax :: VarSet -> Maybe Int
lookupMax :: VarSet -> Maybe Int
lookupMax (VS# Word#
w) =
  if Int# -> Bool
isTrue# (Word#
w Word# -> Word# -> Int#
`eqWord#` Word#
0##) then
    Maybe Int
forall a. Maybe a
Nothing
  else
    Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# Word#
w)))
lookupMax (VB# ByteArray#
bs) =
  -- We know that the highest bit must be in the final word of 'bs',
  -- as there are no leading 0 words in our varsets.
  let len :: Int#
len = ByteArray# -> Int#
wordArraySize# ByteArray#
bs
  in Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
highestBitWord# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs (Int#
len Int# -> Int# -> Int#
-# Int#
1#)))))

-- | The number of entries in the variable set.
size :: VarSet -> Int
size :: VarSet -> Int
size (VS# Word#
w) = Int# -> Int
I# (Word# -> Int#
word2Int# (Word# -> Word#
popCnt# Word#
w))
size (VB# ByteArray#
bs) = Int# -> Int
I# (Word# -> Int#
word2Int# (ByteArray# -> Word#
bigNatPopCount# ByteArray#
bs))
{-# INLINABLE size #-}

-- | Check if two variable sets are disjoint.
disjoint :: VarSet -> VarSet -> Bool
disjoint :: VarSet -> VarSet -> Bool
disjoint (VS# Word#
w1) (VS# Word#
w2) = Int# -> Bool
isTrue# (Word# -> Word# -> Int#
disjointWord# Word#
w1 Word#
w2)
disjoint VarSet
vs1 VarSet
vs2 = Int# -> Bool
isTrue# (VarSet -> VarSet -> Int#
disjointSlow# VarSet
vs1 VarSet
vs2)
{-# INLINABLE disjoint #-}

-- | Slow path for 'disjoint'.
disjointSlow# :: VarSet -> VarSet -> Int#
disjointSlow# :: VarSet -> VarSet -> Int#
disjointSlow# (VS# Word#
w1) (VS# Word#
w2) = Word# -> Word# -> Int#
disjointWord# Word#
w1 Word#
w2
disjointSlow# (VB# ByteArray#
bs1) (VS# Word#
w2) = Word# -> Word# -> Int#
disjointWord# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs1 Int#
0#) Word#
w2
disjointSlow# (VS# Word#
w1) (VB# ByteArray#
bs2) = Word# -> Word# -> Int#
disjointWord# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs2 Int#
0#) Word#
w1
disjointSlow# (VB# ByteArray#
bs1) (VB# ByteArray#
bs2) = ByteArray# -> ByteArray# -> Int#
byteArrayDisjoint# ByteArray#
bs1 ByteArray#
bs2
{-# NOINLINE disjointSlow# #-}

--------------------------------------------------------------------------------
-- Combining variable sets

-- | Union two variable sets.
union :: VarSet -> VarSet -> VarSet
union :: VarSet -> VarSet -> VarSet
union (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`or#` Word#
w2)
union VarSet
vs1 VarSet
vs2 = VarSet -> VarSet -> VarSet
unionSlow VarSet
vs1 VarSet
vs2
{-# INLINABLE union #-}

-- | Slow path for 'union'.
unionSlow :: VarSet -> VarSet -> VarSet
unionSlow :: VarSet -> VarSet -> VarSet
unionSlow (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`or#` Word#
w2)
unionSlow (VS# Word#
w1) (VB# ByteArray#
bs2) = ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatOrWord# ByteArray#
bs2 Word#
w1)
unionSlow (VB# ByteArray#
bs1) (VS# Word#
w2) = ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatOrWord# ByteArray#
bs1 Word#
w2)
unionSlow (VB# ByteArray#
bs1) (VB# ByteArray#
bs2) = ByteArray# -> VarSet
VB# (ByteArray# -> ByteArray# -> ByteArray#
bigNatOr ByteArray#
bs1 ByteArray#
bs2)
{-# NOINLINE unionSlow #-}

-- | Union a collection of variable sets together.
unions :: P.Foldable f => f VarSet -> VarSet
unions :: forall (f :: * -> *). Foldable f => f VarSet -> VarSet
unions = (VarSet -> VarSet -> VarSet) -> VarSet -> f VarSet -> VarSet
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Fold.foldl' VarSet -> VarSet -> VarSet
union VarSet
empty
{-# SPECIALIZE unions :: [VarSet] -> VarSet #-}

-- | Take the intersection of two variable sets.
intersection :: VarSet -> VarSet -> VarSet
intersection :: VarSet -> VarSet -> VarSet
intersection (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2)
intersection VarSet
vs1 VarSet
vs2 = VarSet -> VarSet -> VarSet
intersectionSlow VarSet
vs1 VarSet
vs2
{-# INLINABLE intersection #-}

-- | Slow path for 'intersection'.
intersectionSlow :: VarSet -> VarSet -> VarSet
intersectionSlow :: VarSet -> VarSet -> VarSet
intersectionSlow (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2)
intersectionSlow (VS# Word#
w1) (VB# ByteArray#
bs2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`and#` ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs2 Int#
0#)
intersectionSlow (VB# ByteArray#
bs1) (VS# Word#
w2) = Word# -> VarSet
VS# (ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs1 Int#
0# Word# -> Word# -> Word#
`and#` Word#
w2)
intersectionSlow (VB# ByteArray#
bs1) (VB# ByteArray#
bs2) = ByteArray# -> VarSet
VB# (ByteArray# -> ByteArray# -> ByteArray#
bigNatAnd ByteArray#
bs1 ByteArray#
bs2)
{-# NOINLINE intersectionSlow #-}

-- | Set difference of two variable sets.
difference :: VarSet -> VarSet -> VarSet
difference :: VarSet -> VarSet -> VarSet
difference (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`andNot#` Word#
w2)
difference VarSet
vs1 VarSet
vs2 = VarSet -> VarSet -> VarSet
differenceSlow VarSet
vs1 VarSet
vs2
{-# INLINABLE difference #-}

-- | Slow path for 'difference'.
differenceSlow :: VarSet -> VarSet -> VarSet
differenceSlow :: VarSet -> VarSet -> VarSet
differenceSlow (VS# Word#
w1) (VS# Word#
w2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`andNot#` Word#
w2)
differenceSlow (VS# Word#
w1) (VB# ByteArray#
bs2) = Word# -> VarSet
VS# (Word#
w1 Word# -> Word# -> Word#
`andNot#` ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs2 Int#
0#)
differenceSlow (VB# ByteArray#
bs1) (VS# Word#
w2) = ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatAndNotWord# ByteArray#
bs1 Word#
w2)
differenceSlow (VB# ByteArray#
bs1) (VB# ByteArray#
bs2) = ByteArray# -> VarSet
byteArrayToVarSet# (ByteArray# -> ByteArray# -> ByteArray#
bigNatAndNot ByteArray#
bs1 ByteArray#
bs2)
{-# NOINLINE differenceSlow #-}

-- | Infix operator for 'difference'.
(\\) :: VarSet -> VarSet -> VarSet
\\ :: VarSet -> VarSet -> VarSet
(\\) = VarSet -> VarSet -> VarSet
difference
{-# INLINE (\\) #-}

-- | Take a relative complement of a variable set.
complement :: Int -> VarSet -> VarSet
complement :: Int -> VarSet -> VarSet
complement Int
n VarSet
vs = Int -> VarSet
full Int
n VarSet -> VarSet -> VarSet
\\ VarSet
vs
{-# INLINE complement #-}

-- | @isSubsetOf vs1 vs2@ determines if @vs1@ is a subset of @vs2@.
isSubsetOf :: VarSet -> VarSet -> Bool
isSubsetOf :: VarSet -> VarSet -> Bool
isSubsetOf (VS# Word#
w1) (VS# Word#
w2) = Int# -> Bool
isTrue# ((Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2) Word# -> Word# -> Int#
`eqWord#` Word#
w1)
isSubsetOf VarSet
vs1 VarSet
vs2 = Int# -> Bool
isTrue# (VarSet -> VarSet -> Int#
isSubsetOfSlow# VarSet
vs1 VarSet
vs2)
{-# INLINABLE isSubsetOf #-}

-- | Slow path for 'isSubsetOf'.
isSubsetOfSlow# :: VarSet -> VarSet -> Int#
isSubsetOfSlow# :: VarSet -> VarSet -> Int#
isSubsetOfSlow# (VS# Word#
w1) (VS# Word#
w2) = (Word#
w1 Word# -> Word# -> Word#
`and#` Word#
w2) Word# -> Word# -> Int#
`eqWord#` Word#
w1
isSubsetOfSlow# (VS# Word#
w1) (VB# ByteArray#
bs2) = (Word#
w1 Word# -> Word# -> Word#
`and#` ByteArray# -> Int# -> Word#
indexWordArray# ByteArray#
bs2 Int#
0#) Word# -> Word# -> Int#
`eqWord#` Word#
w1
isSubsetOfSlow# (VB# ByteArray#
bs1) (VS# Word#
w2) = Int#
0#
isSubsetOfSlow# (VB# ByteArray#
bs1) (VB# ByteArray#
bs2) = ByteArray# -> ByteArray# -> Int#
byteArrayIsSubsetOf# ByteArray#
bs1 ByteArray#
bs2
{-# NOINLINE isSubsetOfSlow# #-}

-- | @inRange lo hi vs@ determines if all the variables in @vs@
-- are within the range @[lo+1 .. hi-1]@.
inRange :: Int -> Int -> VarSet -> Bool
inRange :: Int -> Int -> VarSet -> Bool
inRange Int
lo Int
hi VarSet
vs =
  -- This might seem inefficient, but is optimal
  -- when @range lo hi@ and @vs@ both fit within a single word.
  VarSet
vs VarSet -> VarSet -> Bool
`isSubsetOf` Int -> Int -> VarSet
range Int
lo Int
hi
{-# INLINE inRange #-}

--------------------------------------------------------------------------------
-- Filters

-- | Split a variable set into elements that are strictly smaller than
-- @n@ and elements that are strictly larger.
--
-- Note that this is strict in both components of the pair. If you only
-- need one half, use 'filterLT' or 'filterGT'.
split :: Int -> VarSet -> (VarSet, VarSet)
split :: Int -> VarSet -> (VarSet, VarSet)
split Int
n VarSet
vs =
  let !lo :: VarSet
lo = VarSet -> VarSet -> VarSet
intersection VarSet
vs (Int -> VarSet
full Int
n)
      !hi :: VarSet
hi = VarSet -> VarSet -> VarSet
difference VarSet
vs (Int -> VarSet
full (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  in (VarSet
lo, VarSet
hi)
{-# INLINE split #-}

-- | Filter a variable set to elements that are less than to an index.
filterLT :: Int -> VarSet -> VarSet
filterLT :: Int -> VarSet -> VarSet
filterLT Int
n VarSet
vs = VarSet -> VarSet -> VarSet
intersection VarSet
vs (Int -> VarSet
full Int
n)
{-# INLINE filterLT #-}

-- | Filter a variable set to elements that are less than or equal to an index.
filterLE :: Int -> VarSet -> VarSet
filterLE :: Int -> VarSet -> VarSet
filterLE Int
n VarSet
vs = VarSet -> VarSet -> VarSet
intersection VarSet
vs (Int -> VarSet
full (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
{-# INLINE filterLE #-}

-- | Filter a variable set to elements that are greater than to an index.
filterGT :: Int -> VarSet -> VarSet
filterGT :: Int -> VarSet -> VarSet
filterGT Int
n VarSet
vs = VarSet -> VarSet -> VarSet
difference VarSet
vs (Int -> VarSet
full (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
{-# INLINE filterGT #-}

-- | Filter a variable set to elements that are greater than or equal to an index.
filterGE :: Int -> VarSet -> VarSet
filterGE :: Int -> VarSet -> VarSet
filterGE Int
n VarSet
vs = VarSet -> VarSet -> VarSet
difference VarSet
vs (Int -> VarSet
full Int
n)
{-# INLINE filterGE #-}

--------------------------------------------------------------------------------
-- Folds

-- $varSetFolds
-- Like most unordered containers, @foldr@ and @foldl@ are somewhat ambigious
-- names for folds over a 'VarSet'. For the sake of consistency, we adopt the
-- same convention as 'Data.IntSet.IntSet', where
-- @
-- VarSet.foldr f x = foldr f x . VarSet.toAscList
-- VarSet.foldl f x = foldl f x . VarSet.toAscList
-- @
--
-- Unfortunately, this convention the opposite of how our sets are laid out
-- in memory.

-- | Lazily fold over the elements of a variable set in ascending order.
foldr :: (Int -> a -> a) -> a -> VarSet -> a
foldr :: forall a. (Int -> a -> a) -> a -> VarSet -> a
foldr Int -> a -> a
f a
a (VS# Word#
w) = (Int -> a -> a) -> a -> Word# -> a
forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBits# Int -> a -> a
f a
a Word#
w
foldr Int -> a -> a
f a
a (VB# ByteArray#
bs) = (Int -> a -> a) -> a -> ByteArray# -> a
forall a. (Int -> a -> a) -> a -> ByteArray# -> a
byteArrayFoldrBits# Int -> a -> a
f a
a ByteArray#
bs

-- | Lazily fold over the elements of a variable set in descending order.
--
-- This does not have the same deficiencies 'Data.List.foldl', as we can
-- start the fold at the back of the 'VarSet'.
foldl :: (a -> Int -> a) -> a -> VarSet -> a
foldl :: forall a. (a -> Int -> a) -> a -> VarSet -> a
foldl a -> Int -> a
f a
a (VS# Word#
w) = (a -> Int -> a) -> a -> Word# -> a
forall a. (a -> Int -> a) -> a -> Word# -> a
wordFoldlBits# a -> Int -> a
f a
a Word#
w
foldl a -> Int -> a
f a
a (VB# ByteArray#
bs) = (a -> Int -> a) -> a -> ByteArray# -> a
forall a. (a -> Int -> a) -> a -> ByteArray# -> a
byteArrayFoldlBits# a -> Int -> a
f a
a ByteArray#
bs

-- | Strictly fold over the elements of a variable set in ascending order.
--
-- This does not have the same deficiencies 'Data.List.foldr'', as we can
-- start the fold at the back of the 'VarSet'.
foldr' :: (Int -> a -> a) -> a -> VarSet -> a
foldr' :: forall a. (Int -> a -> a) -> a -> VarSet -> a
foldr' Int -> a -> a
f a
a (VS# Word#
w) = (Int -> a -> a) -> a -> Word# -> a
forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsStrict# Int -> a -> a
f a
a Word#
w
foldr' Int -> a -> a
f a
a (VB# ByteArray#
bs) = (Int -> a -> a) -> a -> ByteArray# -> a
forall a. (Int -> a -> a) -> a -> ByteArray# -> a
byteArrayFoldrBitsStrict# Int -> a -> a
f a
a ByteArray#
bs

-- | Strictly fold over the elements of a variable set in descending order.
foldl' :: (Int -> a -> a) -> a -> VarSet -> a
foldl' :: forall a. (Int -> a -> a) -> a -> VarSet -> a
foldl' Int -> a -> a
f a
a (VS# Word#
w) = (Int -> a -> a) -> a -> Word# -> a
forall a. (Int -> a -> a) -> a -> Word# -> a
wordFoldrBitsStrict# Int -> a -> a
f a
a Word#
w
foldl' Int -> a -> a
f a
a (VB# ByteArray#
bs) = (Int -> a -> a) -> a -> ByteArray# -> a
forall a. (Int -> a -> a) -> a -> ByteArray# -> a
byteArrayFoldrBitsStrict# Int -> a -> a
f a
a ByteArray#
bs


--------------------------------------------------------------------------------
-- Views

-- | Retrieve the smallest index in the variable set along with the set sans that
-- element, or 'Nothing' if the set was empty.
minView :: VarSet -> Maybe (Int, VarSet)
minView :: VarSet -> Maybe (Int, VarSet)
minView VarSet
vs =
  case VarSet -> Maybe Int
lookupMin VarSet
vs of
    Just Int
i ->
      let !vs' :: VarSet
vs' = Int -> VarSet -> VarSet
delete Int
i VarSet
vs
      in (Int, VarSet) -> Maybe (Int, VarSet)
forall a. a -> Maybe a
Just (Int
i, VarSet
vs')
    Maybe Int
Nothing -> Maybe (Int, VarSet)
forall a. Maybe a
Nothing
{-# INLINABLE minView #-}

-- | Retrieve the smallest index in the variable set along with the set sans that
-- element, or 'Nothing' if the set was empty.
maxView :: VarSet -> Maybe (Int, VarSet)
maxView :: VarSet -> Maybe (Int, VarSet)
maxView VarSet
vs =
  case VarSet -> Maybe Int
lookupMax VarSet
vs of
    Just Int
i ->
      let !vs' :: VarSet
vs' = Int -> VarSet -> VarSet
delete Int
i VarSet
vs
      in (Int, VarSet) -> Maybe (Int, VarSet)
forall a. a -> Maybe a
Just (Int
i, VarSet
vs')
    Maybe Int
Nothing -> Maybe (Int, VarSet)
forall a. Maybe a
Nothing
{-# INLINABLE maxView #-}

--------------------------------------------------------------------------------
-- Contextual operations

-- | Remove the first @n@ entries from the variable set,
-- and shift the indices of all other entries down by @n@.
strengthen :: Int -> VarSet -> VarSet
strengthen :: Int -> VarSet -> VarSet
strengthen (I# Int#
n) (VS# Word#
w) =
  -- A strengthening by 64+ will literally never happen, so we are better
  -- off avoiding the branch then we are saving an allocation by
  -- re-using 'empty' when the shift result would be 0.
  Word# -> VarSet
VS# (Word#
w Word# -> Int# -> Word#
`shiftRL#` Int#
n)
strengthen (I# Int#
n) (VB# ByteArray#
bs) =
  ByteArray# -> VarSet
byteArrayToVarSet# (ByteArray# -> Word# -> ByteArray#
bigNatShiftR# ByteArray#
bs (Int# -> Word#
int2Word# Int#
n))

-- | Shift all indices in the variable set up by @n@.
weaken :: Int -> VarSet -> VarSet
weaken :: Int -> VarSet -> VarSet
weaken (I# Int#
n) (VS# Word#
w) =
  if Int# -> Bool
isTrue# (Word# -> Word#
clz# Word#
w Word# -> Word# -> Int#
`geWord#` Int# -> Word#
int2Word# Int#
n) then
    Word# -> VarSet
VS# (Word#
w Word# -> Int# -> Word#
`uncheckedShiftL#` Int#
n)
  else
    ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatShiftL# (Word# -> ByteArray#
bigNatFromWord# Word#
w) (Int# -> Word#
int2Word# Int#
n))
weaken (I# Int#
n) (VB# ByteArray#
bs) =
  ByteArray# -> VarSet
VB# (ByteArray# -> Word# -> ByteArray#
bigNatShiftL# ByteArray#
bs (Int# -> Word#
int2Word# Int#
n))

--------------------------------------------------------------------------------
-- Conversions

-- | Convert a variable set to a descending list of indices.
toDescList :: VarSet -> [Int]
toDescList :: VarSet -> [Int]
toDescList = ([Int] -> Int -> [Int]) -> [Int] -> VarSet -> [Int]
forall a. (a -> Int -> a) -> a -> VarSet -> a
foldl ((Int -> [Int] -> [Int]) -> [Int] -> Int -> [Int]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) []
{-# INLINE toDescList #-}

-- | Convert a variable set to an ascending list of indices.
toAscList :: VarSet -> [Int]
toAscList :: VarSet -> [Int]
toAscList = (Int -> [Int] -> [Int]) -> [Int] -> VarSet -> [Int]
forall a. (Int -> a -> a) -> a -> VarSet -> a
foldr (:) []
{-# INLINE toAscList #-}

--------------------------------------------------------------------------------
-- Debugging

-- | Trace the size of a variable set to the eventlog.
traceVarSetSize :: VarSet -> VarSet
traceVarSetSize :: VarSet -> VarSet
traceVarSetSize vs :: VarSet
vs@(VS# Word#
_) = String -> VarSet -> VarSet
forall a. String -> a -> a
traceEvent (String
"VarSet.size=1") VarSet
vs
traceVarSetSize vs :: VarSet
vs@(VB# ByteArray#
bs) = String -> VarSet -> VarSet
forall a. String -> a -> a
traceEvent (String
"VarSet.size=" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Word -> String
forall a. Show a => a -> String
show (ByteArray# -> Word
bigNatSize ByteArray#
bs)) VarSet
vs