Search code examples
haskellunboxing

Looking for «instance (Enum a, Bounded a) => IArray UArray a»

I'm looking for a way to have Enum a => UArray a (which makes sense to me as we can trivially map enums to Int and back by toEnum and fromEnum)

So far I tried to steal code of UArray Int from Data.Array.Base and smuggle a few toEnums and fromEnums here and there:

{-# LANGUAGE MagicHash, UnboxedTuples #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}

module UArrays where

import           Data.Array.Base
import           Data.Array.ST
import           Data.Array.Unboxed

import           GHC.Base -- (Int(I#), Int#(..))
import           GHC.Prim -- (indexIntArray#, readIntArray#, writeIntArray#)
import           GHC.ST (ST(..), runST)

import           Unsafe.Coerce

instance (Enum a, Bounded a) => IArray UArray a where
    {-# INLINE bounds #-}
    bounds (UArray l u _ _) = (l, u)
    {-# INLINE numElements #-}
    numElements (UArray _ _ n _) = n
    {-# INLINE unsafeArray #-}
    unsafeArray lu ies = runST (unsafeArrayUArray lu ies minBound)
        {-# INLINE unsafeAt #-}
    unsafeAt (UArray _ _ _ arr#) (I# i#) =
        I# $ fromEnum (indexIntArray# arr# i#)
    {-# INLINE unsafeReplace #-}
    unsafeReplace arr ies = runST (unsafeReplaceUArray arr ies)
    {-# INLINE unsafeAccum #-}
    unsafeAccum f arr ies = runST (unsafeAccumUArray f arr ies)
    {-# INLINE unsafeAccumArray #-}
    unsafeAccumArray f initialValue lu ies =
      runST (unsafeAccumArrayUArray f initialValue lu ies)

-- data STUArray s i e = STUArray !i !i !Int (GHC.Prim.MutableByteArray# s)
instance (Enum a, Bounded a) => MArray (STUArray s) a (ST s) where
    {-# INLINE getBounds #-}
    getBounds (STUArray l u _ _) = return (l, u)
    {-# INLINE getNumElements #-}
    getNumElements (STUArray _ _ n _) = return n
    {-# INLINE unsafeNewArray_ #-}
    unsafeNewArray_ (l, u) = unsafeNewArraySTUArray_ (l, u) wORD_SCALE
    {-# INLINE newArray_ #-}
    newArray_ arrBounds = newArray arrBounds minBound
    {-# INLINE unsafeRead #-}
    -- unsafeRead :: GHC.Arr.Ix i => a i e -> Int -> m e
    unsafeRead (STUArray _ _ _ marr#) (I# i#) =
      ST $ \ s1# ->
      case readIntArray# marr# i# s1# of
        (# s2#, e# #) -> (# s2#, I# (toEnum e#) #)
    {-# INLINE unsafeWrite #-}
    -- unsafeWrite :: GHC.Arr.Ix i => a i e -> Int -> e -> m ()
    unsafeWrite (STUArray _ _ _ marr#) (I# i#) (I# e#) =
      ST $ \ s1# ->
      case writeIntArray# marr# (unsafeCoerce i#) (I# $ fromEnum e#) s1# of
        s2# -> (# s2#, () #)

But of course it doesn't compile:

[2 of 4] Compiling UArrays          ( UArrays.hs, interpreted )

UArrays.hs:27:14:
    Couldn't match expected type `Int#' with actual type `Int'
    In the return type of a call of `fromEnum'
    In the second argument of `($)', namely
      `fromEnum (indexIntArray# arr# i#)'
    In the expression: I# $ fromEnum (indexIntArray# arr# i#)

UArrays.hs:52:45:
    Couldn't match expected type `Int' with actual type `Int#'
    In the first argument of `toEnum', namely `e#'
    In the first argument of `I#', namely `(toEnum e#)'
    In the expression: I# (toEnum e#)

UArrays.hs:57:57:
    Couldn't match expected type `Int#' with actual type `Int'
    In the return type of a call of `fromEnum'
    In the second argument of `($)', namely `fromEnum e#'
    In the third argument of `writeIntArray#', namely
      `(I# $ fromEnum e#)'
Failed, modules loaded: Utils.

Also there is no magical unboxInt :: Int -> Int# in GHC.*, and pattern-matching on I# doesn't yield Int but an Int# instead, yet somehow UArray Int exists and works on Int#s.

I have also found a post about making an UArray for newtypes, but it doesn't seem to apply because it relies on unsafeCoerce. I tried that but it made a funny listArray (0, 54) $ cycle [Red, Yellow, Green] in which all constructors were Blue.

Am I on the wrong track?

Update:

It works now, here is the source code:

Solution

  • You need to realize that Int is a boxed integer that is constructed from an unboxed integer Int# via the constructor I#. So:

    -- These functions aren't practical; they just demonstrate the types.
unboxInt :: Int -> Int#
unboxInt (I# unboxed) = unboxed

boxInt :: Int# -> Int
boxInt = I#

    Since fromEnum already returns a boxed integer, you don't have to re-box it by applying I# again, so in this code snippet, for instance:

    {-I# $-} fromEnum (indexIntArray# arr# i#)

    ... you can simply leave out the I# constructor. Similarly, when using toEnum, you should apply the I# constructor to get a boxed integer out of an unboxed integer.

    As @leftaroundabout mentioned, this boxing and unboxing in addition to the complexity that fromEnum and toEnum can have (Especially for tuples, etc) might lead to less performance compared to using simple boxed Arrays.