Playing with Stack Traces when indexing into byte arrays

stack_trace_primitive.hs

{-# language BangPatterns #-}
{-# language MagicHash #-}
{-# language UnboxedTuples #-}

import GHC.Exts
import GHC.IO (IO(..))
import Control.Exception (SomeException,toException)
import System.IO.Error (userError)

main :: IO ()
main = do
  putStrLn "Enter 1 or 2"
  n <- getLine
  case n of
    "1" -> firstCommand
    "2" -> secondCommand
    _ -> error "invalid command"

die :: IO a
die = IO (raiseIO# myException)

myException :: SomeException
myException = toException (userError "Do not use command #2")

firstCommand :: IO ()
firstCommand = do
  marr <- newIntArray 100
  let go !ix (x : xs) = if ix < 100
        then do
          writeIntArray marr ix x
          go (ix + 1) xs
        else pure ()
  go 0 (squareNumbersHelper 0)
  arr <- unsafeFreezeIntArray marr
  putStrLn "Enter a number under 100"
  n <- getLine
  let m = indexIntArray arr (read n)
  print m

secondCommand :: IO ()
secondCommand = die

squareNumbersHelper :: Int -> [Int]
squareNumbersHelper !n =
  (n * n) : squareNumbersHelper (n + 1)

data IntArray = IntArray ByteArray#
data MutableIntArray = MutableIntArray (MutableByteArray# RealWorld)

-- This implementation is only correct on a 64-bit platform
newIntArray :: Int -> IO MutableIntArray
newIntArray (I# n#)
  = IO (\s# ->
      case newByteArray# (n# *# 8# ) s# of
        (# s'#, arr# #) -> (# s'#, MutableIntArray arr# #)
    )

-- | Read a primitive value from the primitive array.
indexIntArray :: IntArray -> Int -> Int
indexIntArray (IntArray arr#) (I# i#) = I# (indexIntArray# arr# i#)

-- | Write an element to the given index.
writeIntArray ::
     MutableIntArray -- ^ array
  -> Int -- ^ index
  -> Int -- ^ element
  -> IO ()
writeIntArray (MutableIntArray arr#) (I# i#) (I# x) = IO $ \s0 ->
  case writeIntArray# arr# i# x s0 of
    s1 -> (# s1, () #)

-- | Convert a mutable byte array to an immutable one without copying. The
-- array should not be modified after the conversion.
unsafeFreezeIntArray :: MutableIntArray -> IO IntArray
unsafeFreezeIntArray (MutableIntArray arr#) = IO $ \s# ->
  case unsafeFreezeByteArray# arr# s# of
    (# s'#, arr'# #) -> (# s'#, IntArray arr'# #)