#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#include "MachDeps.h"
#endif
module Data.CritBit.Types.Internal
(
CritBitKey(..)
, CritBit(..)
, Set(..)
, BitMask
, Node(..)
, foldlWithKey
, foldlWithKey'
, foldrWithKey
, foldrWithKey'
, toList
) where
import Control.DeepSeq (NFData(..))
import Data.Bits (Bits, (.|.), (.&.), shiftL, shiftR)
import Data.ByteString (ByteString)
import Data.Foldable hiding (toList)
import Data.Monoid (Monoid(..))
import Data.Text ()
import Data.Text.Internal (Text(..))
import Data.Word (Word, Word8, Word16, Word32, Word64)
import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as B
import qualified Data.Text.Array as T
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
type BitMask = Word16
data Node k v =
Internal {
ileft, iright :: !(Node k v)
, ibyte :: !Int
, iotherBits :: !BitMask
}
| Leaf k v
| Empty
deriving (Eq, Show)
instance (NFData k, NFData v) => NFData (Node k v) where
rnf (Internal l r _ _) = rnf l `seq` rnf r
rnf (Leaf k v) = rnf k `seq` rnf v
rnf Empty = ()
instance Functor (Node k) where
fmap f i@(Internal l r _ _) = i { ileft = fmap f l, iright = fmap f r }
fmap f (Leaf k v) = Leaf k (f v)
fmap _ Empty = Empty
instance Foldable (Node k) where
foldl f z m = foldlWithKey (\a _ v -> f a v) z (CritBit m)
foldr f z m = foldrWithKey (\_ v a -> f v a) z (CritBit m)
foldMap f (Internal l r _ _) = mappend (foldMap f l) (foldMap f r)
foldMap f (Leaf _ v) = f v
foldMap _ Empty = mempty
foldlWithKey :: (a -> k -> v -> a) -> a -> CritBit k v -> a
foldlWithKey f z m = foldlWithKeyWith (\_ b -> b) f z m
foldlWithKey' :: (a -> k -> v -> a) -> a -> CritBit k v -> a
foldlWithKey' f z m = foldlWithKeyWith seq f z m
foldrWithKey :: (k -> v -> a -> a) -> a -> CritBit k v -> a
foldrWithKey f z m = foldrWithKeyWith (\_ b -> b) f z m
foldrWithKey' :: (k -> v -> a -> a) -> a -> CritBit k v -> a
foldrWithKey' f z m = foldrWithKeyWith seq f z m
foldlWithKeyWith :: (a -> a -> a) -> (a -> k -> v -> a) -> a -> CritBit k v -> a
foldlWithKeyWith maybeSeq f z0 (CritBit root) = go z0 root
where
go z (Internal left right _ _) = let z' = go z left
in z' `maybeSeq` go z' right
go z (Leaf k v) = f z k v
go z Empty = z
foldrWithKeyWith :: (a -> a -> a) -> (k -> v -> a -> a) -> a -> CritBit k v -> a
foldrWithKeyWith maybeSeq f z0 (CritBit root) = go root z0
where
go (Internal left right _ _) z = let z' = go right z
in z' `maybeSeq` go left z'
go (Leaf k v) z = f k v z
go Empty z = z
newtype CritBit k v = CritBit (Node k v)
deriving (Eq, NFData, Functor, Foldable)
instance (Show k, Show v) => Show (CritBit k v) where
show t = "fromList " ++ show (toList t)
class (Eq k) => CritBitKey k where
byteCount :: k -> Int
getByte :: k -> Int -> Word16
instance CritBitKey ByteString where
byteCount = B.length
getByte bs n
| n < B.length bs = fromIntegral (B.unsafeIndex bs n) .|. 256
| otherwise = 0
instance CritBitKey Text where
byteCount (Text _ _ len) = len `shiftL` 1
getByte (Text arr off len) n
| n < len `shiftL` 1 =
let word = T.unsafeIndex arr (off + (n `shiftR` 1))
byteInWord = (word `shiftR` ((n .&. 1) `shiftL` 3)) .&. 0xff
in byteInWord .|. 256
| otherwise = 0
#if WORD_SIZE_IN_BITS == 64
# define WORD_SHIFT 3
#else
# define WORD_SHIFT 2
#endif
instance CritBitKey (U.Vector Word8) where
byteCount = G.length
getByte = getByteV 0
instance CritBitKey (U.Vector Word16) where
byteCount = (`shiftL` 1) . G.length
getByte = getByteV 1
instance CritBitKey (U.Vector Word32) where
byteCount = (`shiftL` 2) . G.length
getByte = getByteV 2
instance CritBitKey (U.Vector Word64) where
byteCount = (`shiftL` 3) . G.length
getByte = getByteV 3
instance CritBitKey (U.Vector Word) where
byteCount = (`shiftL` WORD_SHIFT) . G.length
getByte = getByteV WORD_SHIFT
instance CritBitKey (U.Vector Char) where
byteCount = (`shiftL` 2) . G.length
getByte = getByteV_ fromEnum 2
instance CritBitKey (V.Vector Word8) where
byteCount = G.length
getByte = getByteV 0
instance CritBitKey (V.Vector Word16) where
byteCount = (`shiftL` 1) . G.length
getByte = getByteV 1
instance CritBitKey (V.Vector Word32) where
byteCount = (`shiftL` 2) . G.length
getByte = getByteV 2
instance CritBitKey (V.Vector Word64) where
byteCount = (`shiftL` 3) . G.length
getByte = getByteV 3
instance CritBitKey (V.Vector Word) where
byteCount = (`shiftL` WORD_SHIFT) . G.length
getByte = getByteV WORD_SHIFT
instance CritBitKey (V.Vector Char) where
byteCount = (`shiftL` 2) . G.length
getByte = getByteV_ fromEnum 2
getByteV :: (Bits a, Integral a, G.Vector v a) => Int -> v a -> Int -> Word16
getByteV = getByteV_ id
getByteV_ :: (Bits a, Integral a, G.Vector v b) =>
(b -> a) -> Int -> v b -> Int -> Word16
getByteV_ convert shiftSize = \v n ->
if n < G.length v `shiftL` shiftSize
then reindex shiftSize n $ \wordOffset shiftRight ->
let word = convert (G.unsafeIndex v wordOffset)
byteInWord = (word `shiftR` shiftRight) .&. 255
in fromIntegral byteInWord .|. 256
else 0
reindex :: Int -> Int -> (Int -> Int -> r) -> r
reindex shiftSize n f = f wordOffset shiftRight
where
wordOffset = n `shiftR` shiftSize
shiftRight = (size (n .&. size)) `shiftL` 3
where size = (1 `shiftL` shiftSize) 1
toList :: CritBit k v -> [(k, v)]
toList (CritBit root) = go root []
where
go (Internal l r _ _) next = go l (go r next)
go (Leaf k v) next = (k,v) : next
go Empty next = next
newtype Set a = Set (CritBit a ())
deriving (Eq, NFData)