Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

unionArrayBy: Find next 1-bits with countTrailingZeros #395

Merged
merged 6 commits into from
Mar 31, 2022
Merged

unionArrayBy: Find next 1-bits with countTrailingZeros #395

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a780a8d
unionArrayBy: Find next 1-bits with countTrailingZeros
sjakobi Mar 24, 2022
b00fbd7
Add union.ByteString benchmark
sjakobi Mar 24, 2022
eedc326
Avoid range checks due to bit, clearBit, testBit
sjakobi Mar 24, 2022
605800e
Remove outdated comment
sjakobi Mar 28, 2022
1d386c0
Cleanup
sjakobi Mar 28, 2022
f55e79a
More cleanup
sjakobi Mar 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 22 additions & 19 deletions Data/HashMap/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
{-# LANGUAGE DeriveLift #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
Expand Down Expand Up @@ -143,8 +144,9 @@ import Control.Applicative (Const (..))
import Control.DeepSeq (NFData (..), NFData1 (..), NFData2 (..))
import Control.Monad.ST (ST, runST)
import Data.Bifoldable (Bifoldable (..))
import Data.Bits (complement, popCount, unsafeShiftL,
unsafeShiftR, (.&.), (.|.))
import Data.Bits (bit, clearBit, complement,
countTrailingZeros, popCount, testBit,
unsafeShiftL, unsafeShiftR, (.&.), (.|.))
import Data.Coerce (coerce)
import Data.Data (Constr, Data (..), DataType)
import Data.Functor.Classes (Eq1 (..), Eq2 (..), Ord1 (..), Ord2 (..),
Expand Down Expand Up @@ -1625,23 +1627,24 @@ unionArrayBy f !b1 !b2 !ary1 !ary2 = A.run $ do
let b' = b1 .|. b2
mary <- A.new_ (popCount b')
-- iterate over nonzero bits of b1 .|. b2
-- it would be nice if we could shift m by more than 1 each time
let ba = b1 .&. b2
go !i !i1 !i2 !m
| m > b' = return ()
| b' .&. m == 0 = go i i1 i2 (m `unsafeShiftL` 1)
| ba .&. m /= 0 = do
x1 <- A.indexM ary1 i1
x2 <- A.indexM ary2 i2
A.write mary i $! f x1 x2
go (i+1) (i1+1) (i2+1) (m `unsafeShiftL` 1)
| b1 .&. m /= 0 = do
A.write mary i =<< A.indexM ary1 i1
go (i+1) (i1+1) i2 (m `unsafeShiftL` 1)
| otherwise = do
A.write mary i =<< A.indexM ary2 i2
go (i+1) i1 (i2+1) (m `unsafeShiftL` 1)
go 0 0 0 (b' .&. negate b') -- XXX: b' must be non-zero
let go !b
| b == 0 = return ()
| otherwise = do
let ba = b1 .&. b2
c = countTrailingZeros b
m = bit c
sjakobi marked this conversation as resolved.
Show resolved Hide resolved
i = sparseIndex b' m
i1 = sparseIndex b1 m
i2 = sparseIndex b2 m
sjakobi marked this conversation as resolved.
Show resolved Hide resolved
t <- if | testBit ba c -> do
x1 <- A.indexM ary1 i1
x2 <- A.indexM ary2 i2
return $! f x1 x2
| testBit b1 c -> A.indexM ary1 i1
| otherwise -> A.indexM ary2 i2
A.write mary i t
go (clearBit b c)
go b'
return mary
-- TODO: For the case where b1 .&. b2 == b1, i.e. when one is a
-- subset of the other, we could use a slightly simpler algorithm,
Expand Down