{-|
Small sets represented as a 'Word64'.
For an element type, you have to make sure that the Enum instance maps values to the 0-63 range.
-}

module Utils.SmallSet where

import Utils.List
import Data.Word (Word64)
import Data.Bits

type role SmallSet representational
newtype SmallSet a = SmallSet Word64
  deriving (Eq, Ord) via Word64

instance (Show a, Enum a) => Show (SmallSet a) where
  showsPrec p s = showParen (p > 10) $
    ("SmallSet " ++) . showsPrec 11 (fold (flip (:)) [] s)

instance Semigroup (SmallSet a) where
  {-# inline (<>) #-}
  (<>) = \(SmallSet x) (SmallSet y) -> SmallSet (x .|. y)

instance Monoid (SmallSet a) where
  {-# inline mempty #-}
  mempty = SmallSet 0

{-# inlinable toList #-}
toList :: Enum a => SmallSet a -> List a
toList = fold (flip Cons) Nil

{-# inlinable fromList #-}
fromList :: Enum a => [a] -> SmallSet a
fromList = foldl' (\s a -> insert a s) mempty

{-# inline fold #-}
fold :: Enum a => (b -> a -> b) -> b -> SmallSet a -> b
fold = \f b (SmallSet as) ->
  let go i b 0  = b
      go i b as = case countTrailingZeros as of
        shift -> go (i + shift + 1) (f b (toEnum (i + shift)))
                    (unsafeShiftR as (shift + 1))
  in go 0 b as

{-# inline null #-}
null :: SmallSet a -> Bool
null = \s -> s == mempty

{-# inline insert #-}
insert :: Enum a => a -> SmallSet a -> SmallSet a
insert = \a (SmallSet as) -> SmallSet (as .|. unsafeShiftL 1 (fromEnum a))

{-# inline member #-}
member :: Enum a => a -> SmallSet a -> Bool
member = \a (SmallSet as) ->
  toEnum (fromIntegral (unsafeShiftR as (fromEnum a) .&. 1))

{-# inline delete #-}
delete :: Enum a => a -> SmallSet a -> SmallSet a
delete = \a (SmallSet as) -> SmallSet (as .&. complement (unsafeShiftL 1 (fromEnum a)))

{-# inline notMember #-}
notMember :: Enum a => a -> SmallSet a -> Bool
notMember = \a s -> not (member a s)

{-# inline singleton #-}
singleton :: Enum a => a -> SmallSet a
singleton = \a -> SmallSet (unsafeShiftL 1 (fromEnum a))