-----------------------------------------------------------------------------
-- |
-- Module      :  EnumSet
-- Copyright   :  (c) David F. Place 2006
-- Derived from Data.Set by Daan Leijen
-- License     :  BSD-style
-- Maintainer  :  David F. Place
-- Stability   :  ??
-- Portability :  portable
--
-- An efficient implementation of sets over small enumerations.
--
-- This module is intended to be imported @qualified@, to avoid name
-- clashes with "Prelude" functions.  eg.
--
-- >  import EnumSet as Set
--
-- The implementation of 'EnumSet' is based on bit-wise operations.
-----------------------------------------------------------------------------

module EnumSet 
    (
            -- * Set type
            Set          
            , UniverseSet(universe)
            -- * Operators
            , (\\)

            -- * Query
            , null
            , size
            , member
            , isSubsetOf
            , isProperSubsetOf
            
            -- * Construction
            , empty
            , singleton
            , insert
            , delete
            
            -- * Combine
            , union, unions
            , difference
            , intersection
            , complement
            , complementWith

            -- * Filter
            , filter
            , partition
            , split
            , splitMember

            -- * Map
	    , map
	    , mapMonotonic

            -- * Fold
            , fold

            -- * Min\/Max
            , findMin
            , findMax
            , deleteMin
            , deleteMax
            , deleteFindMin
            , deleteFindMax

            -- * Conversion

            -- ** List
            , elems
            , toList
            , fromList
            
            -- ** Ordered list
            , toAscList
            , fromAscList
            , fromDistinctAscList
)  where
import Prelude hiding (filter,foldr,null,map)
import Data.Bits hiding (complement)
import Data.Word
import Data.List (foldl',intersperse,sort)
import Data.Monoid (Monoid(..))

{--------------------------------------------------------------------
  Operators
--------------------------------------------------------------------}
infixl 9 \\ --

(\\) :: Set a -> Set a -> Set a
m1 \\ m2 = difference m1 m2

{--------------------------------------------------------------------
  Sets are bit strings of width @wordLength@.
--------------------------------------------------------------------}
-- | A set of values @a@.
newtype Set a = Set Word deriving (Eq)

class (Bounded a, Enum a) => UniverseSet a where
    universe :: Set a
    universe = fromList [minBound .. maxBound]

wordLength = foldBits f 0 (maxBound::Word)
    where
      f z _ = z+1
    
check :: String -> Int -> Int 
check msg x  
    | x < wordLength = x
    | otherwise = error $ "EnumSet."++msg++"` beyond word size."


{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}
null :: Set a -> Bool
null (Set 0) = True
null _       = False

size :: Set a -> Int
size (Set w) = foldBits f 0 w
    where
      f z _ = z+1

member :: Enum a => a -> Set a -> Bool
member x (Set w) = testBit w $ fromEnum x

{--------------------------------------------------------------------
  Construction
--------------------------------------------------------------------}
empty :: Set a
empty = Set 0

singleton :: Enum a => a -> Set a
singleton x =
    Set $ setBit 0 $ check "singleton" $ fromEnum x

{--------------------------------------------------------------------
  Insertion, Deletion
--------------------------------------------------------------------}
insert :: Enum a => a -> Set a -> Set a
insert x (Set w) =
    Set $ setBit w $ check "insert" $ fromEnum x

delete :: Enum a => a -> Set a -> Set a
delete x (Set w) = 
    Set $ clearBit w $ fromEnum x

{--------------------------------------------------------------------
  Subset
--------------------------------------------------------------------}
isProperSubsetOf :: Set a -> Set a -> Bool
isProperSubsetOf x y = (x /= y) && (isSubsetOf x y)

isSubsetOf :: Set a -> Set a -> Bool
isSubsetOf x y = (x `union` y) == y

{--------------------------------------------------------------------
  Minimal, Maximal
--------------------------------------------------------------------}
findMin :: Enum a => Set a -> a
findMin (Set w) = toEnum $ findMinIndex w

findMinIndex :: Word -> Int
findMinIndex 0 = 
    error "EnumSet.findMin: empty set has no minimal element"
findMinIndex w = f 0 w
    where 
      f i w
          | 1 == (w .&. 1) = i
          | otherwise = f (i+1) (w `shiftR` 1)

findMax :: Enum a => Set a -> a
findMax (Set w) = toEnum $ findMaxIndex w

findMaxIndex :: Word -> Int
findMaxIndex 0 = 
    error "EnumSet.findMax: empty set has no maximal element"
findMaxIndex w = foldBits (\_ i -> i) 0 w

deleteMin :: Set a -> Set a
deleteMin (Set 0) = empty
deleteMin (Set w) = Set $ clearBit w $ findMinIndex w

deleteMax :: Set a -> Set a
deleteMax (Set 0) = empty
deleteMax (Set w) = Set $ clearBit w $ findMaxIndex w

deleteFindMin :: Enum a => Set a -> (a,Set a)
deleteFindMin s@(Set 0) = 
    (error 
     "EnumSet.deleteFindMin: can not return the minimal element of an empty set", 
     s)
deleteFindMin s = (min,delete min s)
    where min = findMin s

deleteFindMax :: Enum a => Set a -> (a,Set a)
deleteFindMax s@(Set 0) = 
    (error 
     "EnumSet.deleteFindMax: can not return the maximal element of an empty set", 
     s)
deleteFindMax s = (max,delete max s)
    where max = findMax s

{--------------------------------------------------------------------
  Union. 
--------------------------------------------------------------------}
union :: Set a -> Set a -> Set a
union (Set x) (Set y) = Set $ x .|. y

unions :: [Set a] -> Set a
unions = foldl' union empty

{--------------------------------------------------------------------
  Difference
--------------------------------------------------------------------}
difference :: Set a -> Set a -> Set a
difference (Set x) (Set y) = Set $ (x .|. y) `xor` y

{--------------------------------------------------------------------
  Intersection
--------------------------------------------------------------------}
intersection :: Set a -> Set a -> Set a
intersection (Set x) (Set y) = Set $ x .&. y

{--------------------------------------------------------------------
  Complement
--------------------------------------------------------------------}
complement :: UniverseSet a => Set a -> Set a
complement x = complementWith u x
    where u = universe `asTypeOf` x

complementWith :: Set a -> Set a -> Set a
complementWith (Set u) (Set x) = Set $ u `xor` x

{--------------------------------------------------------------------
  Filter and partition
--------------------------------------------------------------------}
filter :: Enum a => (a -> Bool) -> Set a -> Set a
filter p (Set w) = Set $ foldBits f 0 w
    where 
      f z i 
        | p $ toEnum i = setBit z i
        | otherwise = z

partition :: Enum a => (a -> Bool) -> Set a -> (Set a,Set a)
partition p (Set w) = (Set yay,Set nay)
    where 
      (yay,nay) = foldBits f (0,0) w
      f (x,y) i 
          | p $ toEnum i = (setBit x i,y)
          | otherwise = (x,setBit y i)

{----------------------------------------------------------------------
  Map
----------------------------------------------------------------------}
map :: (Enum a,Enum b) => (a -> b) -> Set a -> Set b
map f (Set w) = Set $ foldBits fold 0 w
    where 
      fold z i = setBit z $ check "map" $ fromEnum $ f (toEnum i)

-- @'mapMonotonic'@ is provided for compatibility with the 
-- Data.Set interface.
mapMonotonic :: (Enum a,Enum b) => (a -> b) -> Set a -> Set b
mapMonotonic = map

{--------------------------------------------------------------------
  Fold
--------------------------------------------------------------------}
fold :: Enum a => (b -> a -> b) -> b -> Set a -> b
fold f z (Set w) = foldBits folder z w
    where
      folder z i = f z $ toEnum i

-- Not a post-order fold.  Does it matter?
foldr :: Enum a => (b -> a -> b) -> b -> Set a -> b
foldr = fold

{--------------------------------------------------------------------
  List variations 
--------------------------------------------------------------------}
elems :: Enum a => Set a -> [a]
elems = toList

{--------------------------------------------------------------------
  Lists 
--------------------------------------------------------------------}
toList :: Enum a => Set a -> [a]
toList (Set w) = reverse $ foldBits f [] w
    where
      f z i = (toEnum i) : z

toAscList :: (Ord a,Enum a) => Set a -> [a]
toAscList = sort . toList

fromList :: Enum a => [a] -> Set a
fromList xs = Set $ foldl' f 0 xs
    where 
      f z x = setBit z $ check "fromList" $ fromEnum x

fromAscList :: Enum a => [a] -> Set a
fromAscList = fromList

fromDistinctAscList :: Enum a => [a] -> Set a
fromDistinctAscList = fromList

{--------------------------------------------------------------------
  Show
--------------------------------------------------------------------}
instance (Enum a, Show a) => Show (Set a) where
    show xs = 
        "{"++(concat $ intersperse "," [show x | x <- toList xs])++"}"

{--------------------------------------------------------------------
  Split
--------------------------------------------------------------------}
split :: (Ord a, Enum a) => a -> Set a -> (Set a,Set a)
split x s = (lesser,greater)
    where (lesser,_,greater) = splitMember x s

splitMember :: (Ord a, Enum a) => a -> Set a -> (Set a,Bool,Set a)
splitMember x (Set w) = (Set lesser,isMember,Set greater)
    where
      (lesser,isMember,greater) = foldBits f (0,False,0) w
      f (lesser,isMember,greater) i =
        case compare (toEnum i) x of
          GT -> (lesser,isMember,setBit greater i)
          LT -> (setBit lesser i,isMember,greater)
          EQ -> (lesser,True,greater)

{--------------------------------------------------------------------
  Utility functions. 
--------------------------------------------------------------------}

foldBits :: Bits c => (a -> Int -> a) -> a -> c -> a
foldbits _ z 0  = z
foldBits f z bs = foldBits' f 0 bs z

foldBits' :: Bits c => (a -> Int -> a) -> Int -> c -> a -> a
foldBits' f i bs z
    | bs == 0 = z
    | otherwise = z' `seq` foldBits' f i' bs' z'
    where z' | 1 == bs .&. 1 = f z i
             | otherwise =  z
          i' = i + 1
          bs' = bs `shiftR` 1

{--------------------------------------------------------------------
  Ord 
--------------------------------------------------------------------}
instance (Enum a,Ord a) => Ord (Set a) where
    compare a b = compare (toAscList a) (toAscList b)

{--------------------------------------------------------------------
  Monoid
--------------------------------------------------------------------}
instance Enum a => Monoid (Set a) where
    mempty  = empty
    mappend = union
    mconcat = unions