From 5874a30af459892454cf696f5e6ca913b7f400cc Mon Sep 17 00:00:00 2001
From: Simon Jakobi <simon.jakobi@gmail.com>
Date: Sun, 6 Mar 2022 00:13:33 +0100
Subject: Improve setBit for Natural

Previously the default definition was used, which involved allocating
intermediate Natural values.

Fixes #21173.
---
 libraries/ghc-bignum/src/GHC/Num/BigNat.hs  |  1 +
 libraries/ghc-bignum/src/GHC/Num/Natural.hs | 11 +++++++++++
 2 files changed, 12 insertions(+)

(limited to 'libraries/ghc-bignum/src/GHC')

diff --git a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
index 5fa81b7e5b..bab3a97190 100644
--- a/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/BigNat.hs
@@ -1045,6 +1045,7 @@ bigNatClearBit# a n
 
 -- | BigNat set bit
 bigNatSetBit# :: BigNat# -> Word# -> BigNat#
+{-# NOINLINE bigNatSetBit# #-}
 bigNatSetBit# a n
    -- check the current bit value
    | isTrue# (bigNatTestBit# a n) = a
diff --git a/libraries/ghc-bignum/src/GHC/Num/Natural.hs b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
index 72b646501d..8fe737bf07 100644
--- a/libraries/ghc-bignum/src/GHC/Num/Natural.hs
+++ b/libraries/ghc-bignum/src/GHC/Num/Natural.hs
@@ -421,6 +421,17 @@ naturalBit# i
 naturalBit :: Word -> Natural
 naturalBit (W# i) = naturalBit# i
 
+-- | @since 1.3
+naturalSetBit# :: Natural -> Word# -> Natural
+naturalSetBit# (NS n) i
+  | isTrue# (i `ltWord#` WORD_SIZE_IN_BITS##) = NS (n `or#` (1## `uncheckedShiftL#` word2Int# i))
+  | True                                      = NB (bigNatSetBit# (bigNatFromWord# n) i)
+naturalSetBit# (NB n) i                       = NB (bigNatSetBit# n i)
+
+-- | @since 1.3
+naturalSetBit :: Natural -> Word -> Natural
+naturalSetBit !n (W# i) = naturalSetBit# n i
+
 -- | Compute greatest common divisor.
 naturalGcd :: Natural -> Natural -> Natural
 {-# NOINLINE naturalGcd #-}
-- 
cgit v1.2.1