diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for atomic-modify-generic + +## 0.1.0.0 -- YYYY-mm-dd + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,24 @@ +BSD 2-Clause License + +Copyright (c) 2023, David Feuer + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/atomic-modify-general.cabal b/atomic-modify-general.cabal new file mode 100644 --- /dev/null +++ b/atomic-modify-general.cabal @@ -0,0 +1,78 @@ +cabal-version: 3.0 +name: atomic-modify-general +version: 0.1.0.0 +synopsis: Generalizations of atomicModifyIORef +description: + @base@ provides + + @ + atomicModifyIORef :: IORef a -> (a -> (a, b)) -> IO b + atomicModifyIORef2 :: IORef a -> (a -> (a, b)) -> IO (a, (a, b)) + @ + + to modify the value in an @IORef@ and return a result (and, in the + case of @atomicModifyIORef2@, also return the old value). + + In "Data.IORef.AtomicModify", we generalize this from pairs to arbitrary + types for which the user can provide a function to extract the new + value to store in the @IORef@. + + In "Data.IORef.AtomicModify.Generic", we offer a faster but more restricted + version taking advantage of the fact that the primop used to implement + @atomicModifyIORef2@ actually works for /somewhat/ more general record types + than @atomicModifyIORef2@ accepts. + +homepage: https://github.com/treeowl/atomic-modify-general +license: BSD-2-Clause +license-file: LICENSE +author: David Feuer +maintainer: David.Feuer@gmail.com +copyright: 2023 David Feuer +category: Concurrency +build-type: Simple +extra-doc-files: CHANGELOG.md +-- extra-source-files: +tested-with: + GHC == 8.4.4, GHC == 8.6.5, GHC == 8.8.4, GHC == 8.10.7 + , GHC == 9.0.2, GHC == 9.2.7, GHC == 9.4.4, GHC == 9.6.1 +source-repository head + type: git + location: http://github.com/treeowl/atomic-modify-general.git + +common warnings + ghc-options: -Wall + +library + import: warnings + exposed-modules: + Data.IORef.AtomicModify + , Data.IORef.AtomicModify.Generic + other-modules: + Data.IORef.AtomicModify.SmallArraySize + -- other-extensions: + build-depends: base >= 4.11 && < 4.19 + , primitive + hs-source-dirs: src + if impl (ghc >= 8.10) + hs-source-dirs: compat-sa-size-recent + else + hs-source-dirs: compat-sa-size-old + if impl (ghc < 8.8) + hs-source-dirs: compat-atomic-modify-old + other-modules: + Data.IORef.AtomicModify.Generic.UnsafeToPair + default-language: Haskell2010 + if impl (ghc < 8.6) + ghc-options: -Wno-unticked-promoted-constructors + +test-suite atomic-modify-general-test + import: warnings + default-language: Haskell2010 + -- other-modules: + -- other-extensions: + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + build-depends: + base >= 4.11.0 && < 4.19 + , atomic-modify-general diff --git a/compat-atomic-modify-old/Data/IORef/AtomicModify/Generic/UnsafeToPair.hs b/compat-atomic-modify-old/Data/IORef/AtomicModify/Generic/UnsafeToPair.hs new file mode 100644 --- /dev/null +++ b/compat-atomic-modify-old/Data/IORef/AtomicModify/Generic/UnsafeToPair.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE Unsafe #-} + +{-# OPTIONS_GHC -fno-worker-wrapper #-} + + +-- | The documentation for `unsafeCoerce` firmly warns that it should not be +-- used to coerce between actually different types. But we want to produce +-- selector thunks and we don't want to rely on Generic nonsense inlining away +-- to do so. So ... we use `unsafeCoerce` in an entirely illegitimate way. We +-- use this module to make sure our bogus coercions don't escape and cause +-- trouble in the optimizer. Using `-fno-worker-wrapper` and `NOINLINE` pragmas +-- prevents GHC from producing unfoldings for these hideous functions, so +-- they're quite opaque. However, it still produces demand signatures and arity +-- info, which *tend* to be good to have right. I'm not sure how much they +-- matter in this particular case. +module Data.IORef.AtomicModify.Generic.UnsafeToPair + ( unsafeToPair + , unsafeFromPair + ) where + +import GHC.Exts (Any) +import Unsafe.Coerce (unsafeCoerce) + +-- | Pretend that an arbitrary type is actually a pair whose first +-- component is whatever we want it to be. This is very unsafe! +unsafeToPair :: t -> (a, Any) +unsafeToPair = unsafeCoerce +{-# NOINLINE unsafeToPair #-} + +-- | Coerce (back) from a (fake) pair to the actual type. +unsafeFromPair :: (a, Any) -> t +unsafeFromPair = unsafeCoerce +{-# NOINLINE unsafeFromPair #-} diff --git a/compat-sa-size-old/Data/IORef/AtomicModify/SmallArraySize.hs b/compat-sa-size-old/Data/IORef/AtomicModify/SmallArraySize.hs new file mode 100644 --- /dev/null +++ b/compat-sa-size-old/Data/IORef/AtomicModify/SmallArraySize.hs @@ -0,0 +1,9 @@ +{-# language MagicHash #-} +{-# language UnboxedTuples #-} +module Data.IORef.AtomicModify.SmallArraySize + ( getSizeofSmallMutableArray# + ) where +import GHC.Exts (SmallMutableArray#, State#, RealWorld, Int#, sizeofSmallMutableArray#) + +getSizeofSmallMutableArray# :: SmallMutableArray# s a -> State# RealWorld -> (# State# RealWorld, Int# #) +getSizeofSmallMutableArray# sma s = (# s, sizeofSmallMutableArray# sma #) diff --git a/compat-sa-size-recent/Data/IORef/AtomicModify/SmallArraySize.hs b/compat-sa-size-recent/Data/IORef/AtomicModify/SmallArraySize.hs new file mode 100644 --- /dev/null +++ b/compat-sa-size-recent/Data/IORef/AtomicModify/SmallArraySize.hs @@ -0,0 +1,5 @@ +{-# language MagicHash #-} +module Data.IORef.AtomicModify.SmallArraySize + ( getSizeofSmallMutableArray# + ) where +import GHC.Exts (getSizeofSmallMutableArray#) diff --git a/src/Data/IORef/AtomicModify.hs b/src/Data/IORef/AtomicModify.hs new file mode 100644 --- /dev/null +++ b/src/Data/IORef/AtomicModify.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE UnboxedTuples #-} + +-- | Atomic 'IORef' and array modification operations for more general result +-- types. +module Data.IORef.AtomicModify + ( atomicModifyIORef3General + , atomicModifyArray3General + , atomicModifySmallArray3General + ) where + +import GHC.IORef (IORef (..), newIORef, readIORef) +import GHC.STRef (STRef (..)) +import GHC.Exts ( casMutVar#, MutVar#, RealWorld, readMutVar#, lazy, State# + , writeMutVar#, Int (..), Int#, MutableArray#, readArray# + , casArray#, SmallMutableArray#, readSmallArray#, casSmallArray# + ) +import Data.IORef.AtomicModify.SmallArraySize + ( getSizeofSmallMutableArray# ) +import GHC.IO (IO (..)) +import System.IO.Unsafe (unsafeDupablePerformIO) +import Data.Primitive.Array +import Data.Primitive.SmallArray +import Control.Monad (when) +import Control.Exception (ArrayException (..), throwIO) + +-- | A version of 'GHC.IORef.atomicModifyIORef2' that takes an arbitrary pair +-- of functions. This function will allocate more than 'atomicModifyIORef2', +-- and will tend to take longer to succeed when there is a lot of contention +-- for the 'IORef'. +-- +-- @ +-- atomicModifyIORef2 ref f = do +-- (old, _new, r) <- atomicModifyIORef2General ref fst f +-- pure (old, r) +-- @ +-- +-- If the first function (the \"extraction function\") is a record field +-- selector (e.g., 'snd'), we do our best to make sure the thunk placed in the +-- 'IORef' is a selector thunk, so the garbage collector can drop the rest of +-- the record once the record is forced. In other cases, callers should +-- generally force the returned @new@ value in order to avoid a potential space +-- leak. +-- +-- Conceptually: +-- +-- @ +-- atomicModifyIORef3General ref extract f = do +-- -- Begin atomic block +-- old <- 'readIORef' ref +-- let r = f old +-- new = extract r +-- 'writeIORef' ref new +-- -- End atomic block +-- r `seq` pure (old, new, r) +-- @ +-- +-- where other threads cannot interfere with the operations in the \"atomic block\". +-- In particular, no other thread can write to the 'IORef' between the 'readIORef' +-- and the 'writeIORef' operations. +atomicModifyIORef3General + :: IORef a -> (t -> a) -> (a -> t) -> IO (a, a, t) +atomicModifyIORef3General (IORef (STRef ref)) extract = \f -> do + -- atomicModifyMutVar2# creates a thunk for the result of applying the user + -- function 'f' to the "old" value read from the IORef, and then edits that thunk + -- in a CAS loop. In Haskell land, thunks are immutable, so we can't exactly + -- do that. Instead, we make an IORef, 'holder', to hold the "old" value, and + -- use 'unsafeDupablePerformIO' to create a thunk that will read it and apply + -- 'f' to the result. We /can/ edit the holder IORef in the CAS loop. + -- Note: since casMutVar# introduces a full memory barrier, any thread reading + -- the 'r' thunk from 'ref' will have "seen" the preceding 'writeMutVar' to 'holder', + -- so it won't get the uninitialized value or anything similarly stale. + holder@(IORef (STRef holder#)) <- newIORef uninitialized + let r = unsafeDupablePerformIO (f <$> readIORef holder) + -- I really don't think r is going to inline anyway, but if it does, we + -- could produce an unnecessary space leak. + {-# NOINLINE r #-} + let new = extract r + IO (\s -> case atomicModifyIORef3General' ref holder# new r s of + (# s', old, new', !res #) -> (# s', (old, new', res) #)) +{-# INLINE atomicModifyIORef3General #-} + +atomicModifyIORef3General' + :: MutVar# RealWorld a -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #) +atomicModifyIORef3General' ref holder new r s1 = + case readMutVar# ref s1 of { (# s2, old #) -> + case writeMutVar# holder old s2 of { s3 -> + case casMutVar# ref old new s3 of { (# s4, flag, _ #) -> + case flag of + -- Why the lazy invocations? + -- + -- In the event that 'old' gets forced, unboxed, and reboxed between + -- 'readMutVar#' and 'casMutVar#', the CAS will never succeed. I doubt + -- that could happen anyway, but let's be sure. + -- + -- If 'r' is forced before the holder is written, it will read an + -- uninitialized value and throw an exception. Ouch. Let's make sure that + -- doesn't happen either. + 0# -> (# s4, lazy old, lazy new, lazy r #) + _ -> atomicModifyIORef3General' ref holder new r s4 }}} +{-# NOINLINE atomicModifyIORef3General' #-} + +uninitialized :: a +uninitialized = error "Uninitialized. This is a bug in atomic-modify-generics." +{-# NOINLINE uninitialized #-} + +-- | A version of 'atomicModifyIORef3General' for 'Array's. See the +-- documentation there. Indexing is performed safely. +atomicModifyArray3General + :: MutableArray RealWorld a -> Int -> (t -> a) -> (a -> t) -> IO (a, a, t) +-- See atomicModifyIORef3General for implementation comments. +-- Why do we perform safe indexing? This operation is expensive enough +-- that I don't think we really have to worry about the cost of a single +-- bounds check. +atomicModifyArray3General mary@(MutableArray mary#) ix@(I# ix#) extract = \f -> do + -- We use unsigned comparison to make sure ix is non-negative + -- and less than the array size with just one comparison. The LLVM + -- backend is clever enough to produce this from the obvious two-sided + -- check, but last I looked the native code generator wasn't. + let !sz = sizeofMutableArray mary + when ((fromIntegral ix :: Word) >= fromIntegral sz) $ outOfBoundsArray ix sz + holder@(IORef (STRef holder#)) <- newIORef uninitialized + let r = unsafeDupablePerformIO (f <$> readIORef holder) + {-# NOINLINE r #-} + let new = extract r + IO (\s -> case atomicModifyArray3General' mary# ix# holder# new r s of + (# s', old, new', !res #) -> (# s', (old, new', res) #)) +{-# INLINE atomicModifyArray3General #-} + +outOfBoundsArray :: Int -> Int -> IO a +outOfBoundsArray ix sz + | ix < 0 = throwIO . IndexOutOfBounds $ + "\natomicModifyArray3General was passed a negative array index of " ++ + show ix ++ "." + | otherwise = throwIO . IndexOutOfBounds $ + "\natomicModifyArray3General was passed an array index of " ++ + show ix ++ ",\nbut an array of only " ++ show sz ++ " elements." +{-# NOINLINE outOfBoundsArray #-} + +atomicModifyArray3General' + :: MutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #) +atomicModifyArray3General' mary ix holder new r s1 = + case readArray# mary ix s1 of { (# s2, old #) -> + case writeMutVar# holder old s2 of { s3 -> + case casArray# mary ix old new s3 of { (# s4, flag, _ #) -> + case flag of + 0# -> (# s4, lazy old, lazy new, lazy r #) + _ -> atomicModifyArray3General' mary ix holder new r s4 }}} +{-# NOINLINE atomicModifyArray3General' #-} + +-- | A version of 'atomicModifyIORef3General' for 'SmallArray's. See the +-- documentation there. Indexing is performed safely. +atomicModifySmallArray3General + :: SmallMutableArray RealWorld a -> Int -> (t -> a) -> (a -> t) -> IO (a, a, t) +-- See atomicModifyIORef3General for implementation comments. +atomicModifySmallArray3General (SmallMutableArray mary#) ix@(I# ix#) extract = \f -> do + IO (\s -> case getSizeofSmallMutableArray# mary# s of (# s', sz# #) -> (# s', I# sz# #)) + >>= \sz -> when ((fromIntegral ix :: Word) >= fromIntegral sz) (outOfBoundsSmallArray ix sz) + holder@(IORef (STRef holder#)) <- newIORef uninitialized + let r = unsafeDupablePerformIO (f <$> readIORef holder) + {-# NOINLINE r #-} + let new = extract r + IO (\s -> case atomicModifySmallArray3General' mary# ix# holder# new r s of + (# s', old, new', !res #) -> (# s', (old, new', res) #)) +{-# INLINE atomicModifySmallArray3General #-} + +atomicModifySmallArray3General' + :: SmallMutableArray# RealWorld a -> Int# -> MutVar# RealWorld a -> a -> t -> State# RealWorld -> (# State# RealWorld, a, a, t #) +atomicModifySmallArray3General' mary ix holder new r s1 = + case readSmallArray# mary ix s1 of { (# s2, old #) -> + case writeMutVar# holder old s2 of { s3 -> + case casSmallArray# mary ix old new s3 of { (# s4, flag, _ #) -> + case flag of + 0# -> (# s4, lazy old, lazy new, lazy r #) + _ -> atomicModifySmallArray3General' mary ix holder new r s4 }}} +{-# NOINLINE atomicModifySmallArray3General' #-} + +outOfBoundsSmallArray :: Int -> Int -> IO a +outOfBoundsSmallArray ix sz + | ix < 0 = throwIO . IndexOutOfBounds $ + "\natomicModifySmallArray3General was passed a negative array index of " ++ + show ix ++ "." + | otherwise = throwIO . IndexOutOfBounds $ + "\natomicModifySmallArray3General was passed an array index of " ++ + show ix ++ ",\nbut an array of only " ++ show sz ++ " elements." +{-# NOINLINE outOfBoundsSmallArray #-} diff --git a/src/Data/IORef/AtomicModify/Generic.hs b/src/Data/IORef/AtomicModify/Generic.hs new file mode 100644 --- /dev/null +++ b/src/Data/IORef/AtomicModify/Generic.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UndecidableInstances #-} +-- Unsafe in the presence of custom Generic instances. +{-# LANGUAGE Unsafe #-} +{-# LANGUAGE ViewPatterns #-} + +-- | Atomic modification for more general records, using GHC generics to check +-- their suitablility. When applicable, this is faster than the general +-- utilities in "Data.IORef.AtomicModify". +module Data.IORef.AtomicModify.Generic + ( atomicModifyIORef2Native + ) where + +import Data.Kind (Constraint, Type) +import GHC.Generics +import GHC.IORef (IORef (..)) +#if MIN_VERSION_base(4,13,0) +import GHC.STRef (STRef (..)) +import GHC.Exts (atomicModifyMutVar2#) +#else +import Data.IORef.AtomicModify.Generic.UnsafeToPair (unsafeToPair, unsafeFromPair) +import Data.IORef.AtomicModify (atomicModifyIORef3General) +#endif +import GHC.TypeLits +import GHC.IO (IO (..)) + +-- This trickery was stolen from Csongor Kiss. I don't think we want to use +-- Generic itself, because we don't actually need the Generic dictionary. +-- We use a type family rather than a type synonym to support GHC <= 9.2, +-- which throw a "could not calculate" error on the synonym definition. Huh. +type family EnsureGenericData t where + EnsureGenericData t = EnsureGeneric' t (Rep t) + (TypeError ('Text "Could not calculate " :<>: 'ShowType (Rep t) :$$: + 'Text "Is it an instance of " :<>: 'ShowType Generic :<>: 'Text "?")) +type family EnsureGeneric' t (rep :: Type -> Type) err :: Constraint where + EnsureGeneric' t (M1 _ ('MetaData _ _ _ 'True) f) _ = TypeError ('ShowType t :<>: 'Text " is a newtype.") + EnsureGeneric' _ _ _ = () + +-- A generalization of 'GHC.IORef.atomicModifyIORef2' to any datatype (not +-- newtype) with exactly one constructor whose first field +-- +-- 1. has the same type as the value in the 'IORef'. +-- 2. is not unpacked. +-- +-- Note that the non-unpackedness criterion can potentially be susceptible +-- to change by various compiler flags! An invalid use might therefore +-- compile successfully with @-O0@ but produce a type error with @-O@ or +-- @-O2@. This can be exacerbated by @-funbox-strict-fields@. To ensure +-- type checking will succeed with a record whose first field is strict and +-- monomorphic, it is best to use @{-# NOUNPACK #-}@ explicitly. +-- +-- This function uses the 'atomicModifyMutVar2#' primop, and will therefore be +-- faster than the more general +-- 'Data.IORef.AtomicModify.atomicModifyIORef2General', and more likely to +-- succeed in a reasonable amount of time when there is substantial contention +-- for the 'IORef'. +-- +-- == WARNING +-- +-- This function is safe when used with /derived/ 'Generic' instances. It may +-- be /very unsafe/ when used with hand-written ones. In particular, we use the +-- type's 'Generic' instance (solely) to determine whether its physical layout +-- is suitable for our purposes. We also rely on certain details of how GHC +-- represents values in memory, and particularly the fact that records are +-- represented using a consistent \"pointers-first\" layout with the first +-- field appearing first. +atomicModifyIORef2Native + :: (EnsureGenericData t, FirstField t (Rep t) ~ a) => IORef a -> (a -> t) -> IO (a, t) +#if MIN_VERSION_base(4,13,0) +atomicModifyIORef2Native (IORef (STRef ref)) f = IO $ \s -> + case atomicModifyMutVar2# ref f s of + (# s', old, !r #) -> (# s', (old, r) #) +#else +atomicModifyIORef2Native ref f = do + -- We don't use fst here because it doesn't inline properly in this context + -- with old GHC versions. + (old, _new, unsafeFromPair -> !r) <- atomicModifyIORef3General ref (\(a, _) -> a) (unsafeToPair . f) + pure (old, r) +#endif + +type family FirstField t rep where + FirstField t (M1 _ ('MetaSel _ _ _ 'DecidedUnpack) f) = + TypeError ('Text "The first field of " :<>: 'ShowType t :<>: 'Text " is unpacked") + FirstField t (M1 i c f) = FirstField t f + FirstField t (_ :+: _) = TypeError ('ShowType t :<>: 'Text " is not a record type") + FirstField t (f :*: _) = FirstField t f + FirstField _ (K1 i c) = c + FirstField t V1 = TypeError ('ShowType t :<>: 'Text " has no constructors") + FirstField t U1 = TypeError ('ShowType t :<>: 'Text " has no fields") diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,14 @@ +{-# language DeriveGeneric #-} +module Main (main) where +import Data.IORef +import Data.IORef.AtomicModify.Generic +import Control.Monad +import GHC.Generics + +data Pair a b = Pair a b + deriving (Show, Generic) + +main :: IO () +main = do + ref <- newIORef (1 :: Int) + replicateM_ 5 $ atomicModifyIORef2Native ref (\i -> Pair (i + 1) i) >>= print