Generate straightline code for inline array allocation

GHC has an optimization for allocating arrays when the size is
statically known -- it'll generate the code allocating and initializing
the array inline (instead of a call to a procedure from
`rts/PrimOps.cmm`).

However, the generated code uses a loop to do the initialization. Since
we already check that the requested size is small (we check against
`maxInlineAllocSize`), we can generate faster straightline code instead.
This brings about 15% improvement for `newSmallArray#` in my testing and
slightly simplifies the code in GHC.

Signed-off-by: Michal Terepeta <michal.terepeta@gmail.com>

3 files changed, 137 insertions, 0 deletions

diff --git a/testsuite/tests/codeGen/should_run/NewSmallArray.hs b/testsuite/tests/codeGen/should_run/NewSmallArray.hs
new file mode 100644
index 0000000000..6ee92ad109
--- /dev/null
+++ b/testsuite/tests/codeGen/should_run/NewSmallArray.hs
@@ -0,0 +1,96 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+-- Tests for creating and initializing a @SmallArray#@ including the
+-- optimiziation where GHC inlines the code instead of calling the
+-- @newSmallArray#@ primop if the length is small enough and known at compile
+-- time.
+module Main where
+
+import GHC.Exts
+import GHC.ST
+
+import Control.Monad (forM_)
+
+
+main :: IO ()
+main = do
+    let !a00 = newSmallArrayWith42 0
+        !a01 = newSmallArrayWith42 1
+        !a02 = newSmallArrayWith42 2
+        !a03 = newSmallArrayWith42 3
+        !a04 = newSmallArrayWith42 4
+        !a05 = newSmallArrayWith42 5
+        !a06 = newSmallArrayWith42 6
+        !a07 = newSmallArrayWith42 7
+        !a08 = newSmallArrayWith42 8
+        !a09 = newSmallArrayWith42 9
+        !a10 = newSmallArrayWith42 10
+        !a11 = newSmallArrayWith42 11
+        !a12 = newSmallArrayWith42 12
+        !a13 = newSmallArrayWith42 13
+        !a14 = newSmallArrayWith42 14
+        !a15 = newSmallArrayWith42 15
+        !a16 = newSmallArrayWith42 16
+        !a17 = newSmallArrayWith42 17
+        !a18 = newSmallArrayWith42 18
+        !a19 = newSmallArrayWith42 19
+        !a20 = newSmallArrayWith42 20
+        !a21 = newSmallArrayWith42 21
+        !a22 = newSmallArrayWith42 22
+        !a23 = newSmallArrayWith42 23
+        !a24 = newSmallArrayWith42 24
+        !a25 = newSmallArrayWith42 25
+        !a26 = newSmallArrayWith42 26
+        !a27 = newSmallArrayWith42 27
+        !a28 = newSmallArrayWith42 28
+        !a29 = newSmallArrayWith42 29
+        !a30 = newSmallArrayWith42 30
+        !a31 = newSmallArrayWith42 31
+        !a32 = newSmallArrayWith42 32
+        !a33 = newSmallArrayWith42 33
+        !a34 = newSmallArrayWith42 34
+        !a35 = newSmallArrayWith42 35
+        !a36 = newSmallArrayWith42 36
+        !a37 = newSmallArrayWith42 37
+        !a38 = newSmallArrayWith42 38
+        !a39 = newSmallArrayWith42 39
+        !all = [ a00, a01, a02, a03, a04, a05, a06, a07, a08, a09
+               , a10, a11, a12, a13, a14, a15, a16, a17, a18, a19
+               , a20, a21, a22, a23, a24, a25, a26, a27, a28, a29
+               , a30, a31, a32, a33, a34, a35, a36, a37, a38, a39
+               ]
+    forM_ all (print . toListArray)
+
+
+data Array a = Array { unArray :: SmallArray# a }
+
+newSmallArrayWith42 :: Int -> Array Int
+newSmallArrayWith42 n = (runST (newArray n 42))
+-- inline to make sure the length is known at compile time
+{-# INLINE newSmallArrayWith42 #-}
+
+newArray :: Int -> a -> ST s (Array a)
+newArray (I# n#) a = ST $ \s1# -> case newSmallArray# n# a s1# of
+    (# s2#, marr# #) -> case unsafeFreezeSmallArray# marr# s2# of
+        (# s3#, arr# #) -> (# s3#, Array arr# #)
+-- inline to make sure the length is known at compile time
+{-# INLINE newArray #-}
+
+toListArray :: Array a -> [a]
+toListArray arr = go 0
+  where
+    go i | i >= lengthArray arr = []
+         | otherwise = indexArray arr i : go (i+1)
+
+indexArray :: Array a -> Int -> a
+indexArray arr i@(I# i#)
+  | i < 0 || i >= len =
+      error $ "bounds error, offset " ++ show i ++ ", length " ++ show len
+  | otherwise = case indexSmallArray# (unArray arr) i# of
+      (# a #) -> a
+  where len = lengthArray arr
+
+lengthArray :: Array a -> Int
+lengthArray arr = I# (sizeofSmallArray# (unArray arr))
diff --git a/testsuite/tests/codeGen/should_run/NewSmallArray.stdout b/testsuite/tests/codeGen/should_run/NewSmallArray.stdout
new file mode 100644
index 0000000000..8833ccc027
--- /dev/null
+++ b/testsuite/tests/codeGen/should_run/NewSmallArray.stdout
@@ -0,0 +1,40 @@
+[]
+[42]
+[42,42]
+[42,42,42]
+[42,42,42,42]
+[42,42,42,42,42]
+[42,42,42,42,42,42]
+[42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
+[42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42,42]
diff --git a/testsuite/tests/codeGen/should_run/all.T b/testsuite/tests/codeGen/should_run/all.T
index c6a6b2736a..189fb72e7f 100644
--- a/testsuite/tests/codeGen/should_run/all.T
+++ b/testsuite/tests/codeGen/should_run/all.T
@@ -137,6 +137,7 @@ test('StaticByteArraySize', normal, compile_and_run, ['-O2'])
 test('CopySmallArray', normal, compile_and_run, [''])
 test('CopySmallArrayStressTest', reqlib('random'), compile_and_run, [''])
 test('SizeOfSmallArray', normal, compile_and_run, [''])
+test('NewSmallArray', normal, compile_and_run, [''])
 test('T9001', normal, compile_and_run, [''])
 test('T9013', omit_ways(['ghci']),  # ghci doesn't support unboxed tuples
      compile_and_run, [''])