{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
module Data.Map.Word8
( Map
, lookup
, null
, size
, empty
, singleton
, union
, unionWith
, insert
, insertWith
, foldrWithKeys
, foldl'
, traverse_
, toList
, fromList
) where
import Prelude hiding (lookup, null)
import Control.Monad.ST.Run (runSmallArrayST)
import Data.Bits (bit, popCount, testBit, unsafeShiftR, (.&.), (.|.))
import Data.Primitive (SmallArray)
import Data.WideWord (Word256)
import Data.Word (Word8)
import qualified Data.Foldable as F
import qualified Data.Primitive as PM
data Map a
= Map
{-# UNPACK #-} !Word256
{-# UNPACK #-} !(SmallArray a)
deriving stock instance (Eq a) => Eq (Map a)
deriving stock instance Functor Map
instance (Show a) => Show (Map a) where
showsPrec :: Int -> Map a -> ShowS
showsPrec Int
p Map a
m = Int -> [(Word8, a)] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (Map a -> [(Word8, a)]
forall a. Map a -> [(Word8, a)]
toList Map a
m)
instance (Semigroup a) => Semigroup (Map a) where
<> :: Map a -> Map a -> Map a
(<>) = (a -> a -> a) -> Map a -> Map a -> Map a
forall a. (a -> a -> a) -> Map a -> Map a -> Map a
unionWith a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)
instance (Semigroup a) => Monoid (Map a) where
mempty :: Map a
mempty = Map a
forall a. Map a
empty
singleton :: Word8 -> a -> Map a
singleton :: forall a. Word8 -> a -> Map a
singleton !Word8
k a
v =
Word256 -> SmallArray a -> Map a
forall a. Word256 -> SmallArray a -> Map a
Map
(Int -> Word256
forall a. Bits a => Int -> a
bit (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Int Word8
k))
((forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s (SmallArray a)) -> SmallArray a
runSmallArrayST (Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
PM.newSmallArray Int
1 a
v ST s (SmallMutableArray s a)
-> (SmallMutableArray s a -> ST s (SmallArray a))
-> ST s (SmallArray a)
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallMutableArray s a -> ST s (SmallArray a)
SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
PM.unsafeFreezeSmallArray))
null :: Map a -> Bool
null :: forall a. Map a -> Bool
null Map a
m = Map a -> Int
forall a. Map a -> Int
size Map a
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
size :: Map a -> Int
size :: forall a. Map a -> Int
size (Map Word256
keys SmallArray a
_) = Word256 -> Int
forall a. Bits a => a -> Int
popCount Word256
keys
empty :: Map a
empty :: forall a. Map a
empty = Word256 -> SmallArray a -> Map a
forall a. Word256 -> SmallArray a -> Map a
Map Word256
0 SmallArray a
forall a. Monoid a => a
mempty
lookup :: Word8 -> Map a -> Maybe a
lookup :: forall a. Word8 -> Map a -> Maybe a
lookup Word8
kw (Map Word256
keys SmallArray a
vals) = case Word256 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word256
keys Int
k of
Bool
False -> Maybe a
forall a. Maybe a
Nothing
Bool
True -> case Int
k of
Int
0 -> a -> Maybe a
forall a. a -> Maybe a
Just (SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
PM.indexSmallArray SmallArray a
vals Int
0)
Int
_ ->
let ix :: Int
ix = Word256 -> Int
forall a. Bits a => a -> Int
popCount (Word256 -> Int -> Word256
forall a. Bits a => a -> Int -> a
unsafeShiftR Word256
forall a. Bounded a => a
maxBound (Int
256 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) Word256 -> Word256 -> Word256
forall a. Bits a => a -> a -> a
.&. Word256
keys)
in a -> Maybe a
forall a. a -> Maybe a
Just (SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
PM.indexSmallArray SmallArray a
vals Int
ix)
where
k :: Int
k = forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word8 @Int Word8
kw
union :: Map a -> Map a -> Map a
union :: forall a. Map a -> Map a -> Map a
union !ma :: Map a
ma@(Map Word256
ksA SmallArray a
vsA) !mb :: Map a
mb@(Map Word256
ksB SmallArray a
vsB)
| Word256
ksA Word256 -> Word256 -> Bool
forall a. Eq a => a -> a -> Bool
== Word256
0 = Map a
mb
| Word256
ksB Word256 -> Word256 -> Bool
forall a. Eq a => a -> a -> Bool
== Word256
0 = Map a
ma
| Bool
otherwise = Word256 -> SmallArray a -> Map a
forall a. Word256 -> SmallArray a -> Map a
Map Word256
ks (SmallArray a -> Map a) -> SmallArray a -> Map a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s (SmallArray a)) -> SmallArray a
runSmallArrayST do
let sz :: Int
sz = Word256 -> Int
forall a. Bits a => a -> Int
popCount Word256
ks
SmallMutableArray s a
dst <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
PM.newSmallArray Int
sz (a -> ST s (SmallMutableArray s a))
-> ST s a -> ST s (SmallMutableArray s a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsA Int
0
((Int, Int, Int) -> Bool -> Bool -> ST s (Int, Int, Int))
-> (Int, Int, Int) -> Word256 -> Word256 -> ST s ()
forall (m :: * -> *) a.
Monad m =>
(a -> Bool -> Bool -> m a) -> a -> Word256 -> Word256 -> m ()
foldlZipBits256
( \(!Int
ix, !Int
ixA, !Int
ixB) Bool
a Bool
b -> case Bool
a of
Bool
True -> do
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PM.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst Int
ix (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsA Int
ixA
(Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, if Bool
b then Int
ixB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 else Int
ixB)
Bool
False -> case Bool
b of
Bool
True -> do
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PM.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst Int
ix (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsB Int
ixB
(Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixA, Int
ixB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
Bool
False -> (Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Int
ixA, Int
ixB)
)
(Int
0, Int
0, Int
0)
Word256
ksA
Word256
ksB
SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
PM.unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst
where
ks :: Word256
ks = Word256
ksA Word256 -> Word256 -> Word256
forall a. Bits a => a -> a -> a
.|. Word256
ksB
unionWith :: (a -> a -> a) -> Map a -> Map a -> Map a
unionWith :: forall a. (a -> a -> a) -> Map a -> Map a -> Map a
unionWith a -> a -> a
g !ma :: Map a
ma@(Map Word256
ksA SmallArray a
vsA) !mb :: Map a
mb@(Map Word256
ksB SmallArray a
vsB)
| Word256
ksA Word256 -> Word256 -> Bool
forall a. Eq a => a -> a -> Bool
== Word256
0 = Map a
mb
| Word256
ksB Word256 -> Word256 -> Bool
forall a. Eq a => a -> a -> Bool
== Word256
0 = Map a
ma
| Bool
otherwise = Word256 -> SmallArray a -> Map a
forall a. Word256 -> SmallArray a -> Map a
Map Word256
ks (SmallArray a -> Map a) -> SmallArray a -> Map a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s (SmallArray a)) -> SmallArray a
runSmallArrayST do
let sz :: Int
sz = Word256 -> Int
forall a. Bits a => a -> Int
popCount Word256
ks
SmallMutableArray s a
dst <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
PM.newSmallArray Int
sz (a -> ST s (SmallMutableArray s a))
-> ST s a -> ST s (SmallMutableArray s a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsA Int
0
((Int, Int, Int) -> Bool -> Bool -> ST s (Int, Int, Int))
-> (Int, Int, Int) -> Word256 -> Word256 -> ST s ()
forall (m :: * -> *) a.
Monad m =>
(a -> Bool -> Bool -> m a) -> a -> Word256 -> Word256 -> m ()
foldlZipBits256
( \(!Int
ix, !Int
ixA, !Int
ixB) Bool
a Bool
b -> case Bool
a of
Bool
True -> case Bool
b of
Bool
True -> do
a
a' <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsA Int
ixA
a
b' <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsB Int
ixB
let !c :: a
c = a -> a -> a
g a
a' a
b'
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PM.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst Int
ix a
c
(Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
Bool
False -> do
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PM.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst Int
ix (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsA Int
ixA
(Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixB)
Bool
False -> case Bool
b of
Bool
True -> do
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PM.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst Int
ix (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
PM.indexSmallArrayM SmallArray a
vsB Int
ixB
(Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
ixA, Int
ixB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
Bool
False -> (Int, Int, Int) -> ST s (Int, Int, Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
ix, Int
ixA, Int
ixB)
)
(Int
0, Int
0, Int
0)
Word256
ksA
Word256
ksB
SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
PM.unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
dst
where
ks :: Word256
ks = Word256
ksA Word256 -> Word256 -> Word256
forall a. Bits a => a -> a -> a
.|. Word256
ksB
insert :: Word8 -> a -> Map a -> Map a
insert :: forall a. Word8 -> a -> Map a -> Map a
insert = (a -> a -> a) -> Word8 -> a -> Map a -> Map a
forall a. (a -> a -> a) -> Word8 -> a -> Map a -> Map a
insertWith a -> a -> a
forall a b. a -> b -> a
const
insertWith :: (a -> a -> a) -> Word8 -> a -> Map a -> Map a
insertWith :: forall a. (a -> a -> a) -> Word8 -> a -> Map a -> Map a
insertWith a -> a -> a
f Word8
k a
v Map a
m = (a -> a -> a) -> Map a -> Map a -> Map a
forall a. (a -> a -> a) -> Map a -> Map a -> Map a
unionWith a -> a -> a
f (Word8 -> a -> Map a
forall a. Word8 -> a -> Map a
singleton Word8
k a
v) Map a
m
foldlZipBits256 ::
(Monad m) =>
(a -> Bool -> Bool -> m a) ->
a ->
Word256 ->
Word256 ->
m ()
foldlZipBits256 :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool -> Bool -> m a) -> a -> Word256 -> Word256 -> m ()
foldlZipBits256 a -> Bool -> Bool -> m a
g !a
a0 !Word256
x !Word256
y = Int -> a -> m ()
go Int
0 a
a0
where
go :: Int -> a -> m ()
go !Int
ix !a
a = case Int
ix of
Int
256 -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Int
_ -> do
let xval :: Bool
xval = Word256 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word256
x Int
ix
let yval :: Bool
yval = Word256 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word256
y Int
ix
a
a' <- a -> Bool -> Bool -> m a
g a
a Bool
xval Bool
yval
Int -> a -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
a'
_foldrBits256 :: (Word8 -> b -> b) -> b -> Word256 -> b
_foldrBits256 :: forall b. (Word8 -> b -> b) -> b -> Word256 -> b
_foldrBits256 Word8 -> b -> b
g b
b0 Word256
w = Int -> b
go Int
0
where
go :: Int -> b
go Int
ix = case Int
ix of
Int
256 -> b
b0
Int
_ -> case Word256 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word256
w Int
ix of
Bool
True -> Word8 -> b -> b
g (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word8 Int
ix) (Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
Bool
False -> Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
foldrWithKeys :: (Word8 -> a -> b -> b) -> b -> Map a -> b
foldrWithKeys :: forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
foldrWithKeys Word8 -> a -> b -> b
g b
b0 (Map Word256
ks SmallArray a
vs) = Int -> Int -> b
go Int
0 Int
0
where
go :: Int -> Int -> b
go !Int
ix !Int
ixVal = case Int
ix of
Int
256 -> b
b0
Int
_ -> case Word256 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word256
ks Int
ix of
Bool
True ->
Word8 -> a -> b -> b
g
(forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Word8 Int
ix)
(SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
PM.indexSmallArray SmallArray a
vs Int
ixVal)
(Int -> Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
ixVal Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
Bool
False -> Int -> Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
ixVal
foldl' :: (b -> a -> b) -> b -> Map a -> b
{-# INLINE foldl' #-}
foldl' :: forall b a. (b -> a -> b) -> b -> Map a -> b
foldl' b -> a -> b
f b
b0 (Map Word256
_ SmallArray a
vs) = (b -> a -> b) -> b -> SmallArray a -> b
forall b a. (b -> a -> b) -> b -> SmallArray a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' b -> a -> b
f b
b0 SmallArray a
vs
traverse_ :: (Applicative m) => (a -> m b) -> Map a -> m ()
{-# INLINE traverse_ #-}
traverse_ :: forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Map a -> m ()
traverse_ a -> m b
f (Map Word256
_ SmallArray a
vs) = (a -> m b) -> SmallArray a -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
F.traverse_ a -> m b
f SmallArray a
vs
toList :: Map a -> [(Word8, a)]
toList :: forall a. Map a -> [(Word8, a)]
toList = (Word8 -> a -> [(Word8, a)] -> [(Word8, a)])
-> [(Word8, a)] -> Map a -> [(Word8, a)]
forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
foldrWithKeys (\Word8
k a
v [(Word8, a)]
b -> (Word8
k, a
v) (Word8, a) -> [(Word8, a)] -> [(Word8, a)]
forall a. a -> [a] -> [a]
: [(Word8, a)]
b) []
fromList :: [(Word8, a)] -> Map a
fromList :: forall a. [(Word8, a)] -> Map a
fromList = (Map a -> (Word8, a) -> Map a) -> Map a -> [(Word8, a)] -> Map a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (\Map a
acc (Word8
k, a
v) -> Map a -> Map a -> Map a
forall a. Map a -> Map a -> Map a
union Map a
acc (Word8 -> a -> Map a
forall a. Word8 -> a -> Map a
singleton Word8
k a
v)) Map a
forall a. Map a
empty