Improve clearBit and complementBit for Natural

Also optimize bigNatComplementBit#.

Fixes #21175, #21181, #21194.

3 files changed, 41 insertions, 4 deletions

diff --git a/libraries/ghc-bignum/changelog.md b/libraries/ghc-bignum/changelog.md
index d57a6d8a68..6816783dca 100644
--- a/libraries/ghc-bignum/changelog.md
+++ b/libraries/ghc-bignum/changelog.md
@@ -3,7 +3,7 @@
 ## 1.3
 
 - Expose backendName
-- Add `naturalSetBit[#]` (#21173)
+- Add `naturalSetBit[#]` (#21173), `naturalClearBit[#]` (#21175), `naturalComplementBit[#]` (#21181)
 
 ## 1.2
 
diff --git a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
index bab3a97190..efe78f10c9 100644
--- a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
@@ -1070,9 +1070,24 @@ bigNatSetBit# a n
 
 -- | Reverse the given bit
 bigNatComplementBit# :: BigNat# -> Word# -> BigNat#
-bigNatComplementBit# bn i
-  | isTrue# (bigNatTestBit# bn i) = bigNatClearBit# bn i
-  | True                          = bigNatSetBit#   bn i
+bigNatComplementBit# a n =
+   let
+      !sz = wordArraySize# a
+      !nw = word2Int# (n `uncheckedShiftRL#` WORD_SIZE_BITS_SHIFT#)
+      !nb = word2Int# (n `and#` WORD_SIZE_BITS_MASK##)
+      d   = nw +# 1# -# sz
+   in if
+      -- result BigNat will have more limbs
+      | isTrue# (d ># 0#)
+      -> withNewWordArray# (nw +# 1#) \mwa s ->
+            case mwaArrayCopy# mwa 0# a 0# sz s of
+               s' -> case mwaFill# mwa 0## (int2Word# sz) (int2Word# (d -# 1#)) s' of
+                  s'' -> writeWordArray# mwa nw (bitW# nb) s''
+
+      | nv <- bigNatIndex# a nw `xor#` bitW# nb
+      -> withNewWordArrayTrimmed# sz \mwa s ->
+            case mwaArrayCopy# mwa 0# a 0# sz s of
+               s' -> writeWordArray# mwa nw nv s'
 
 -------------------------------------------------
 -- Log operations
diff --git a/libraries/ghc-bignum/src/GHC/Num/Natural.hs b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
index 8fe737bf07..3650e7f42b 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Natural.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
@@ -432,6 +432,28 @@ naturalSetBit# (NB n) i                       = NB (bigNatSetBit# n i)
 naturalSetBit :: Natural -> Word -> Natural
 naturalSetBit !n (W# i) = naturalSetBit# n i
 
+-- | @since 1.3
+naturalClearBit# :: Natural -> Word# -> Natural
+naturalClearBit# x@(NS n) i
+  | isTrue# (i `ltWord#` WORD_SIZE_IN_BITS##) = NS (n `and#` not# (1## `uncheckedShiftL#` word2Int# i))
+  | True                                      = x
+naturalClearBit# (NB n) i                     = naturalFromBigNat# (bigNatClearBit# n i)
+
+-- | @since 1.3
+naturalClearBit :: Natural -> Word -> Natural
+naturalClearBit !n (W# i) = naturalClearBit# n i
+
+-- | @since 1.3
+naturalComplementBit# :: Natural -> Word# -> Natural
+naturalComplementBit# (NS n) i
+  | isTrue# (i `ltWord#` WORD_SIZE_IN_BITS##) = NS (n `xor#` (1## `uncheckedShiftL#` word2Int# i))
+  | True                                      = NB (bigNatSetBit# (bigNatFromWord# n) i)
+naturalComplementBit# (NB n) i                = naturalFromBigNat# (bigNatComplementBit# n i)
+
+-- | @since 1.3
+naturalComplementBit :: Natural -> Word -> Natural
+naturalComplementBit !n (W# i) = naturalComplementBit# n i
+
 -- | Compute greatest common divisor.
 naturalGcd :: Natural -> Natural -> Natural
 {-# NOINLINE naturalGcd #-}