1 files changed, 56 insertions, 29 deletions

diff --git a/compiler/GHC/Data/FastMutInt.hs b/compiler/GHC/Data/FastMutInt.hs
index 3b2725d927..5502fbe16f 100644
--- a/compiler/GHC/Data/FastMutInt.hs
+++ b/compiler/GHC/Data/FastMutInt.hs
@@ -14,12 +14,14 @@ module GHC.Data.FastMutInt
   , FastMutInt2
   , newFastMutInt
   , newFastMutInt2
+  , readFirstFastMutInt
+  , readSecondFastMutInt
+  , writeFirstFastMutInt
+  , writeSecondFastMutInt
+  , atomicFetchAddFirstFastMut
+  , atomicFetchAddSecondFastMut
   , readFastMutInt
-  , readFastMutInt2
   , writeFastMutInt
-  , writeFastMutInt2
-  , atomicFetchAddFastMut
-  , atomicFetchAddFastMut2
   )
 where
 
@@ -27,68 +29,93 @@ import GHC.Prelude.Basic
 
 import GHC.Base
 
-data FastMutInts (n :: S) = FastMutInts !(MutableByteArray# RealWorld)
-data S = S1 | S2
+data FastMutInts (n :: VarCount) = FastMutInts !(MutableByteArray# RealWorld)
+-- It's likely possible to generalise this to n-variables, but no
+-- use cases exist so far in GHC, so we currently choose the simplicity
+-- of implementation.
+data VarCount = OneVar | TwoVars
 
-type FastMutInt  = FastMutInts 'S1
-type FastMutInt2 = FastMutInts 'S2
+type FastMutInt  = FastMutInts 'OneVar
+type FastMutInt2 = FastMutInts 'TwoVars
 
+
+-- Keep the old names around for Haddock:
+readFastMutInt :: FastMutInt -> IO Int
+readFastMutInt = readFirstFastMutInt
+
+writeFastMutInt :: FastMutInt -> Int -> IO ()
+writeFastMutInt = writeFirstFastMutInt
+
+
+-- | Allocate a single mutable int with an initial value.
 newFastMutInt :: Int -> IO FastMutInt
 {-# INLINE newFastMutInt #-}
 newFastMutInt n = do
     let size = finiteBitSize (0 :: Int) `unsafeShiftR` 3
     x <- createFastMutInt size
-    writeFastMutInt x n
+    writeFirstFastMutInt x n
     return x
 
+-- | Allocate a pair of mutable ints with initial values.
 newFastMutInt2 :: Int -> Int -> IO FastMutInt2
 {-# INLINE newFastMutInt2 #-}
 newFastMutInt2 n0 n1 = do
     let size = finiteBitSize (0 :: Int) `unsafeShiftR` 2
                   -- only "shiftR 2" to account for "times 2"
     x <- createFastMutInt size
-    writeFastMutInt x n0
-    writeFastMutInt2 x n1
+    writeFirstFastMutInt x n0
+    writeSecondFastMutInt x n1
     return x
 
+-- | Allocate space for n mutable ints.
 createFastMutInt :: Int -> IO (FastMutInts n)
 {-# INLINE createFastMutInt #-}
 createFastMutInt (I# size) = IO $ \s ->
   case newByteArray# size s of
     (# s, arr #) -> (# s, FastMutInts arr #)
 
-readFastMutInt :: FastMutInts n -> IO Int
-{-# INLINE readFastMutInt #-}
-readFastMutInt (FastMutInts arr) = IO $ \s ->
+-- | Read the first int from either a single or pair of
+-- mutable ints.
+readFirstFastMutInt :: FastMutInts n -> IO Int
+{-# INLINE readFirstFastMutInt #-}
+readFirstFastMutInt (FastMutInts arr) = IO $ \s ->
   case readIntArray# arr 0# s of
     (# s, i #) -> (# s, I# i #)
 
-readFastMutInt2 :: FastMutInt2 -> IO Int
-{-# INLINE readFastMutInt2 #-}
-readFastMutInt2 (FastMutInts arr) = IO $ \s ->
+-- | Read the second int from a pair of mutable ints.
+readSecondFastMutInt :: FastMutInt2 -> IO Int
+{-# INLINE readSecondFastMutInt #-}
+readSecondFastMutInt (FastMutInts arr) = IO $ \s ->
   case readIntArray# arr 1# s of
     (# s, i #) -> (# s, I# i #)
 
-writeFastMutInt :: FastMutInts n -> Int -> IO ()
-{-# INLINE writeFastMutInt #-}
-writeFastMutInt (FastMutInts arr) (I# i) = IO $ \s ->
+-- | Write to a single mutable int, or the first slot of
+-- a pair of mutable ints.
+writeFirstFastMutInt :: FastMutInts n -> Int -> IO ()
+{-# INLINE writeFirstFastMutInt #-}
+writeFirstFastMutInt (FastMutInts arr) (I# i) = IO $ \s ->
   case writeIntArray# arr 0# i s of
     s -> (# s, () #)
 
-writeFastMutInt2 :: FastMutInt2 -> Int -> IO ()
-{-# INLINE writeFastMutInt2 #-}
-writeFastMutInt2 (FastMutInts arr) (I# i) = IO $ \s ->
+-- | Write to the second slot of a pair of mutable ints.
+writeSecondFastMutInt :: FastMutInt2 -> Int -> IO ()
+{-# INLINE writeSecondFastMutInt #-}
+writeSecondFastMutInt (FastMutInts arr) (I# i) = IO $ \s ->
   case writeIntArray# arr 1# i s of
     s -> (# s, () #)
 
-atomicFetchAddFastMut :: FastMutInts n -> Int -> IO Int
-{-# INLINE atomicFetchAddFastMut #-}
-atomicFetchAddFastMut (FastMutInts arr) (I# i) = IO $ \s ->
+-- | Atomically modify a single mutable int, or the first slot
+-- of a pair of mutable ints, by the given value.
+atomicFetchAddFirstFastMut :: FastMutInts n -> Int -> IO Int
+{-# INLINE atomicFetchAddFirstFastMut #-}
+atomicFetchAddFirstFastMut (FastMutInts arr) (I# i) = IO $ \s ->
   case fetchAddIntArray# arr 0# i s of
     (# s, n #) -> (# s, I# n #)
 
-atomicFetchAddFastMut2 :: FastMutInt2 -> Int -> IO Int
-{-# INLINE atomicFetchAddFastMut2 #-}
-atomicFetchAddFastMut2 (FastMutInts arr) (I# i) = IO $ \s ->
+-- | Atomically modify the second slot of a pair of mutable ints
+-- by the given value.
+atomicFetchAddSecondFastMut :: FastMutInt2 -> Int -> IO Int
+{-# INLINE atomicFetchAddSecondFastMut #-}
+atomicFetchAddSecondFastMut (FastMutInts arr) (I# i) = IO $ \s ->
   case fetchAddIntArray# arr 1# i s of
     (# s, n #) -> (# s, I# n #)