Cursed Mixed Array

Array.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedDatatypes #-}

module Main where

import GHC.Exts
import GHC.IO
import Unsafe.Coerce
import Data.Kind

data SA = SA (SmallMutableArray# RealWorld Any)


mkArray :: Int -> a -> IO (SA)
mkArray (I# n) initial = IO $ \s ->
    case unsafeCoerce# (newSmallArray# n initial s) of
        (# s', arr #) -> (# s', SA arr #)

readLifted :: SA -> Int -> IO a
readLifted (SA arr) (I# i) = IO (\s ->
    unsafeCoerce# (readSmallArray# arr i s)
    )

data UWrap (a :: UnliftedType) = UWrap a
-- UWrap is just here because we can't return unlifted types in IO
-- If you don't need your result in IO you can eliminate this indirection.
readUnlifted :: forall a. SA -> Int -> IO (UWrap a)
readUnlifted (SA arr) (I# i) = IO (\s ->
    case unsafeCoerce# (readSmallArray# arr i s) of
        (# s', a :: a #) -> (# s', UWrap a #)
    )

writeLifted :: a -> Int -> SA -> IO ()
writeLifted x (I# i) (SA arr) = IO $ \s ->
    case writeSmallArray# (unsafeCoerce# arr) i x s of
        s -> (# s, () #)

writeUnlifted :: (a :: UnliftedType) -> Int -> SA -> IO ()
writeUnlifted x (I# i) (SA arr) = IO $ \s ->
    case writeSmallArray# arr i (unsafeCoerce# x) s of
        s -> (# s, () #)

type UB :: UnliftedType
data UB = UT | UF

showU :: UWrap UB -> String
showU (UWrap UT) = "UT"
showU (UWrap UF) = "UF"

main :: IO ()
main = do
    arr <- mkArray 4 ()
    writeLifted True 0 arr
    writeLifted False 1 arr
    writeUnlifted UT 2 arr
    writeUnlifted UT 3 arr

    (readLifted arr 0 :: IO Bool) >>= print
    (readLifted arr 1 :: IO Bool) >>= print
    (readUnlifted arr 2 :: IO (UWrap UB)) >>= (putStrLn . showU)
    (readUnlifted arr 3 :: IO (UWrap UB)) >>= (putStrLn . showU)

    return ()