From 3dd86a3bd3496cfc342cfa247b77bfbf1e7d1808 Mon Sep 17 00:00:00 2001
From: Bodigrim <andrew.lelechenko@gmail.com>
Date: Sat, 1 Apr 2023 11:20:08 +0100
Subject: [PATCH 1/2] Implement Control.Concurrent.Counter.Unlifted in CMM

---
 Counter.cmm                                 | 96 +++++++++++++++++++++
 atomic-counter.cabal                        |  4 +
 src/Control/Concurrent/Counter/Lifted/ST.hs |  2 +-
 src/Control/Concurrent/Counter/Unlifted.hs  | 66 ++++++++++++--
 4 files changed, 161 insertions(+), 7 deletions(-)
 create mode 100644 Counter.cmm

diff --git a/Counter.cmm b/Counter.cmm
new file mode 100644
index 0000000..89aee32
--- /dev/null
+++ b/Counter.cmm
@@ -0,0 +1,96 @@
+#include "Cmm.h"
+
+INFO_TABLE(stg_Counter, 0, 1, MUT_PRIM, "Counter", "Counter") ()
+{
+  foreign "C" barf("stg_Counter entered!", NULL) never returns;
+}
+
+stg_newCounterzh (W_ x)
+{
+  P_ c;
+  ("ptr" c) = ccall allocate(MyCapability() "ptr", BYTES_TO_WDS(SIZEOF_StgHeader + WDS(1)));
+  SET_HDR(c, stg_Counter_info, CCCS);
+  W_[c + SIZEOF_StgHeader] = x;
+  return (c);
+}
+
+stg_atomicGetCounterzh (P_ c)
+{
+  W_ x;
+  // load_seqcst64 is available since GHC 9.4
+  (x) = prim %load_seqcst64(c + SIZEOF_StgHeader);
+  return (x);
+}
+
+stg_atomicSetCounterzh (P_ c, W_ x)
+{
+  // store_seqcst64 is available since GHC 9.4
+  prim %store_seqcst64(c + SIZEOF_StgHeader, x);
+  return ();
+}
+
+stg_atomicAddCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_add64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_add64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
+
+stg_atomicSubCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_sub64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_sub64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
+
+stg_atomicAndCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_and64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_and64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
+
+stg_atomicOrCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_or64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_or64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
+
+stg_atomicXorCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_xor64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_xor64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
+
+stg_atomicNandCounterzh (P_ c, W_ x)
+{
+  W_ y;
+#if __GLASGOW_HASKELL__ >= 907
+  (y) = prim %fetch_nand64(c + SIZEOF_StgHeader, x);
+#else
+  (y) = ccall hs_atomic_nand64(c + SIZEOF_StgHeader, x);
+#endif
+  return (y);
+}
diff --git a/atomic-counter.cabal b/atomic-counter.cabal
index 7421b4e..31a6a0c 100644
--- a/atomic-counter.cabal
+++ b/atomic-counter.cabal
@@ -78,6 +78,10 @@ library
     src
   build-depends:
     , base >= 4.14 && < 5
+  if impl(ghc >= 9.4) && !arch(javascript)
+    cmm-sources: Counter.cmm
+    ghc-options: -dcmm-lint
+    cpp-options: -DUSE_CMM
 
 library test-utils
   import: ghc-options
diff --git a/src/Control/Concurrent/Counter/Lifted/ST.hs b/src/Control/Concurrent/Counter/Lifted/ST.hs
index ba304b9..7a1573a 100644
--- a/src/Control/Concurrent/Counter/Lifted/ST.hs
+++ b/src/Control/Concurrent/Counter/Lifted/ST.hs
@@ -75,7 +75,7 @@ set
   -> Int
   -> ST s ()
 set (Counter c) (I# x) = ST $ \s1 -> case Unlifted.set c x s1 of
-  s2 -> (# s2, () #)
+  (# s2 #) -> (# s2, () #)
 
 
 {-# INLINE add #-}
diff --git a/src/Control/Concurrent/Counter/Unlifted.hs b/src/Control/Concurrent/Counter/Unlifted.hs
index 51958ce..b9f1435 100644
--- a/src/Control/Concurrent/Counter/Unlifted.hs
+++ b/src/Control/Concurrent/Counter/Unlifted.hs
@@ -16,10 +16,13 @@
 -- 'Control.Concurrent.Counter.Counter' module.
 ----------------------------------------------------------------------------
 
-{-# LANGUAGE CPP              #-}
-{-# LANGUAGE MagicHash        #-}
-{-# LANGUAGE UnboxedTuples    #-}
-{-# LANGUAGE UnliftedNewtypes #-}
+{-# LANGUAGE CPP                  #-}
+{-# LANGUAGE GHCForeignImportPrim #-}
+{-# LANGUAGE KindSignatures       #-}
+{-# LANGUAGE MagicHash            #-}
+{-# LANGUAGE UnboxedTuples        #-}
+{-# LANGUAGE UnliftedFFITypes     #-}
+{-# LANGUAGE UnliftedNewtypes     #-}
 
 module Control.Concurrent.Counter.Unlifted
   ( Counter
@@ -56,6 +59,54 @@ import GHC.Exts
 
 #define ADD_HASH(x) x#
 
+#if defined(USE_CMM) && SIZEOF_HSINT == 8
+
+-- | Memory location that supports select few atomic operations.
+newtype Counter s = Counter (Any :: UnliftedType)
+
+-- | Create new counter with initial value.
+foreign import prim "stg_newCounterzh"
+  new :: Int# -> State# s -> (# State# s, Counter s #)
+
+-- | Atomically read the counter's value.
+foreign import prim "stg_atomicGetCounterzh"
+  get :: Counter s -> State# s -> (# State# s, Int# #)
+
+-- | Atomically assign new value to the counter.
+foreign import prim "stg_atomicSetCounterzh"
+  set :: Counter s -> Int# -> State# s -> (# State# s #)
+
+-- | Atomically add an amount to the counter and return its old value.
+foreign import prim "stg_atomicAddCounterzh"
+  add :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Atomically subtract an amount from the counter and return its old value.
+foreign import prim "stg_atomicSubCounterzh"
+  sub :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Atomically combine old value with a new one via bitwise and. Returns old counter value.
+foreign import prim "stg_atomicAndCounterzh"
+  and :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Atomically combine old value with a new one via bitwise or. Returns old counter value.
+foreign import prim "stg_atomicOrCounterzh"
+  or :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Atomically combine old value with a new one via bitwise xor. Returns old counter value.
+foreign import prim "stg_atomicXorCounterzh"
+  xor :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Atomically combine old value with a new one via bitwise nand. Returns old counter value.
+foreign import prim "stg_atomicNandCounterzh"
+  nand :: Counter s -> Int# -> State# s -> (# State# s, Int# #)
+
+-- | Compare the underlying pointers of two counters.
+sameCounter :: Counter s -> Counter s -> Bool
+sameCounter (Counter x) (Counter y) =
+  isTrue# (reallyUnsafePtrEquality# x y)
+
+#else
+
 -- | Memory location that supports select few atomic operations.
 newtype Counter s = Counter (MutableByteArray# s)
 
@@ -75,8 +126,9 @@ get (Counter arr) = atomicReadIntArray# arr 0#
 
 {-# INLINE set #-}
 -- | Atomically assign new value to the counter.
-set :: Counter s -> Int# -> State# s -> State# s
-set (Counter arr) = atomicWriteIntArray# arr 0#
+set :: Counter s -> Int# -> State# s -> (# State# s #)
+set (Counter arr) n = \s1 -> case atomicWriteIntArray# arr 0# n s1 of
+  s2 -> (# s2 #)
 
 
 {-# INLINE add #-}
@@ -115,3 +167,5 @@ nand (Counter arr) = fetchNandIntArray# arr 0#
 sameCounter :: Counter s -> Counter s -> Bool
 sameCounter (Counter x) (Counter y) =
   isTrue# (sameMutableByteArray# x y)
+
+#endif

From 38f856756e95a864d959ef44fe86982431cc7a88 Mon Sep 17 00:00:00 2001
From: Bodigrim <andrew.lelechenko@gmail.com>
Date: Sat, 15 Apr 2023 00:15:34 +0100
Subject: [PATCH 2/2] Run short benchmarks on CI

---
 .github/workflows/haskell-ci.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/haskell-ci.yaml b/.github/workflows/haskell-ci.yaml
index ba0831c..f83219d 100644
--- a/.github/workflows/haskell-ci.yaml
+++ b/.github/workflows/haskell-ci.yaml
@@ -48,6 +48,7 @@ jobs:
         cd atomic-counter-*/
         cabal build all --enable-tests
         cabal test --enable-tests --test-show-details=direct all
+        cabal bench --enable-benchmarks --benchmark-options='--stdev 100 --timeout 100' all
     - name: Haddock
       run: |
         cd atomic-counter-*/