fixed

commit 6e5599b37430eb80eb15d588a6cef12ba31ba29e
parent a685da885259f2fa2d7bd5ca310268cdc4199321
Author: Jared Tobin <jared@jtobin.io>
Date:   Thu, 23 Jan 2025 19:20:30 +0400

lib: mod

Diffstat:
M
bench/Main.hs
 | 
44
++++++++++++++++++++++++++++++++++++--------
M
lib/Data/Word/Extended.hs
 | 
31
++++++++++++++++++++++++++++++-
M
test/Main.hs
 | 
10
+++++++++-

3 files changed, 75 insertions(+), 10 deletions(-)
diff --git a/bench/Main.hs b/bench/Main.hs
@@ -7,7 +7,7 @@ import Data.Bits ((.|.), (.&.), (.^.))
 import qualified Data.Word.Extended as W
 import Control.DeepSeq
 import Criterion.Main
-import Prelude hiding (or, and, div)
+import Prelude hiding (or, and, div, mod)
 import qualified Prelude (div)
 
 instance NFData W.Word256
@@ -116,9 +116,43 @@ div = bench "div" $ nf (W.div w0) w1 where
   !w1 = W.to_word256
     0x066bd4c3c10e30260cb6e7832af25f15527b089b258a1fef13b6eec3ce73bf06
 
+div_baseline_small :: Benchmark
+div_baseline_small =
+    bench "div, small (baseline)" $ nf (Prelude.div w0) w1
+  where
+    w0, w1 :: Integer
+    !w0 = 0x7fffffed
+    !w1 = 0x7ffbffed
+
+div_small :: Benchmark
+div_small = bench "div, small" $ nf (W.div w0) w1 where
+  !w0 = W.to_word256 0x7fffffed
+  !w1 = W.to_word256 0x7ffbffed
+
+mod_baseline :: Benchmark
+mod_baseline = bench "mod (baseline)" $ nf (Prelude.rem w0) w1 where
+  w0, w1 :: Integer
+  !w0 = 0x41cf50c7d0d65afabcf5ba37750dba71c7db29ec9f20a216d3ef013a59b9188a
+  !w1 = 0x066bd4c3c10e30260cb6e7832af25f15527b089b258a1fef13b6eec3ce73bf06
+
+mod :: Benchmark
+mod = bench "mod" $ nf (W.mod w0) w1 where
+  !w0 = W.to_word256
+    0x41cf50c7d0d65afabcf5ba37750dba71c7db29ec9f20a216d3ef013a59b9188a
+  !w1 = W.to_word256
+    0x066bd4c3c10e30260cb6e7832af25f15527b089b258a1fef13b6eec3ce73bf06
+
 main :: IO ()
 main = defaultMain [
-    or_baseline
+    mul_baseline
+  , mul
+  , div_baseline
+  , div
+  , mod_baseline
+  , mod
+  , div_baseline_small
+  , div_small
+  , or_baseline
   , or
   , and_baseline
   , and
@@ -128,11 +162,5 @@ main = defaultMain [
   , add
   , sub_baseline
   , sub
-  , mul_baseline
-  , mul
-  , mul128_baseline
-  , mul128
-  , div_baseline
-  , div
   ]
 
diff --git a/lib/Data/Word/Extended.hs b/lib/Data/Word/Extended.hs
@@ -12,7 +12,7 @@ import qualified Data.Bits as B
 import qualified Data.Primitive.PrimArray as PA
 import Data.Word (Word64)
 import GHC.Generics
-import Prelude hiding (div)
+import Prelude hiding (div, mod)
 
 fi :: (Integral a, Num b) => a -> b
 fi = fromIntegral
@@ -572,3 +572,32 @@ div a@(Word256 a0 a1 a2 a3) b@(Word256 b0 b1 b2 b3)
       z3 <- PA.readPrimArray quo 3
       pure (Word256 z0 z1 z2 z3)
 
+mod :: Word256 -> Word256 -> Word256
+mod a@(Word256 a0 a1 a2 a3) b@(Word256 b0 b1 b2 b3)
+  | is_zero b || a == b = zero -- ?
+  | a `lt` b = a
+  | is_word64 a = Word256 (a0 `Prelude.rem` b0) 0 0 0
+  | otherwise = runST $ do
+      quo <- PA.newPrimArray 4
+      PA.setPrimArray quo 0 4 0
+      mx <- PA.newPrimArray 4
+      my <- PA.newPrimArray 4
+      PA.writePrimArray mx 0 a0
+      PA.writePrimArray mx 1 a1
+      PA.writePrimArray mx 2 a2
+      PA.writePrimArray mx 3 a3
+      PA.writePrimArray my 0 b0
+      PA.writePrimArray my 1 b1
+      PA.writePrimArray my 2 b2
+      PA.writePrimArray my 3 b3
+      x <- PA.unsafeFreezePrimArray mx
+      y <- PA.unsafeFreezePrimArray my
+      re <- PA.newPrimArray 4
+      PA.setPrimArray re 0 4 0
+      quotrem quo x y (Just re)
+      z0 <- PA.readPrimArray re 0
+      z1 <- PA.readPrimArray re 1
+      z2 <- PA.readPrimArray re 2
+      z3 <- PA.readPrimArray re 3
+      pure (Word256 z0 z1 z2 z3)
+
diff --git a/test/Main.hs b/test/Main.hs
@@ -9,7 +9,7 @@ import qualified Data.Bits as B
 import qualified Data.Primitive.PrimArray as PA
 import Data.Word (Word64)
 import Data.Word.Extended
-import Prelude hiding (and, or, div)
+import Prelude hiding (and, or, div, mod)
 import qualified Prelude (div)
 import Test.Tasty
 import qualified Test.Tasty.HUnit as H
@@ -155,6 +155,12 @@ div_matches (DivMonotonic (a, b)) =
       !rite = to_word256 (a `Prelude.div` b)
   in  left == rite
 
+mod_matches :: DivMonotonic -> Bool
+mod_matches (DivMonotonic (a, b)) =
+  let !left = to_word256 a `mod` to_word256 b
+      !rite = to_word256 (a `rem` b)
+  in  left == rite
+
 -- assertions ------------------------------------------------------------------
 
 quotrem_r_case0 :: H.Assertion
@@ -338,6 +344,8 @@ arithmetic = testGroup "arithmetic" [
       Q.withMaxSuccess 1000 mul_512_matches
   , Q.testProperty "division matches" $
       Q.withMaxSuccess 1000 div_matches
+  , Q.testProperty "mod matches" $
+      Q.withMaxSuccess 1000 mod_matches
   ]
 
 utils :: TestTree