Throw IOError when allocaBytesAligned gets non-power-of-two align

2 files changed, 25 insertions, 2 deletions

diff --git a/libraries/base/Foreign/Marshal/Alloc.hs b/libraries/base/Foreign/Marshal/Alloc.hs
index 6bdf4feadf..76398b80a6 100644
--- a/libraries/base/Foreign/Marshal/Alloc.hs
+++ b/libraries/base/Foreign/Marshal/Alloc.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples,
-             ScopedTypeVariables #-}
+             ScopedTypeVariables, BangPatterns #-}
 
 -----------------------------------------------------------------------------
 -- |
@@ -60,12 +60,15 @@ module Foreign.Marshal.Alloc (
   finalizerFree
 ) where
 
+import Data.Bits                ( Bits, (.&.) )
 import Data.Maybe
 import Foreign.C.Types          ( CSize(..) )
 import Foreign.Storable         ( Storable(sizeOf,alignment) )
 import Foreign.ForeignPtr       ( FinalizerPtr )
 import GHC.IO.Exception
+import GHC.Num
 import GHC.Real
+import GHC.Show
 import GHC.Ptr
 import GHC.Base
 
@@ -142,7 +145,22 @@ allocaBytes (I# size) action = IO $ \ s0 ->
 -- exception), so the pointer passed to @f@ must /not/ be used after this.
 --
 allocaBytesAligned :: Int -> Int -> (Ptr a -> IO b) -> IO b
-allocaBytesAligned (I# size) (I# align) action = IO $ \ s0 ->
+allocaBytesAligned !_size !align !_action
+    | not $ isPowerOfTwo align =
+      ioError $
+        IOError Nothing InvalidArgument
+          "allocaBytesAligned"
+          ("alignment (="++show align++") must be a power of two!")
+          Nothing Nothing
+  where
+    isPowerOfTwo :: (Bits i, Integral i) => i -> Bool
+    isPowerOfTwo x = x .&. (x-1) == 0
+allocaBytesAligned !size !align !action =
+    allocaBytesAlignedAndUnchecked size align action
+{-# INLINABLE allocaBytesAligned #-}
+
+allocaBytesAlignedAndUnchecked :: Int -> Int -> (Ptr a -> IO b) -> IO b
+allocaBytesAlignedAndUnchecked (I# size) (I# align) action = IO $ \ s0 ->
      case newAlignedPinnedByteArray# size align s0 of { (# s1, mbarr# #) ->
      case unsafeFreezeByteArray# mbarr# s1 of { (# s2, barr#  #) ->
      let addr = Ptr (byteArrayContents# barr#) in
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index 02df43857a..06633d9b07 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -20,6 +20,11 @@
 
   * Remove `Data.Semigroup.Option` and the accompanying `option` function.
 
+  * Make `allocaBytesAligned` and `alloca` throw an IOError when the
+    alignment is not a power-of-two. The underlying primop
+    `newAlignedPinnedByteArray#` actually always assumed this but we didn't
+    document this fact in the user facing API until now.
+
 ## 4.15.0.0 *TBA*
 
   * `openFile` now calls the `open` system call with an `interruptible` FFI