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

1 files changed, 20 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