2011-12-15

Unboxed arrays in Haskell

There is a great idea of Data.Ix implemented for nice and fast Data.Array.Unboxed and others. That allows to form arrays, matrices and other schemas (even triangles) of data access. To gain that flexibility we separate storage implementation (i.e. continous sequence of bytes) from indexing schema implementation.
But sometimes we need a way to access more complex objects and still have that magic Unboxed storage. I.e. instead of accessing single object we need two of them. Once again mighty of Haskell type system comes to save us. Just add another index axis and hide it in the same way how Ix do.
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies, DatatypeContexts #-}

import Control.Arrow ((***))
import Data.Array.Base
import Data.Array.IArray
import Data.Array.Unboxed

class Pack t where
    type Item t
    packSize :: t -> Int
    itemsUnpack :: t -> [Item t]
    packItems :: [Item t] -> t

instance Pack (a, a) where
    type Item (a, a) = a
    packSize _ = 2
    itemsUnpack (x, y) = [x, y]
    packItems [x, y] = (x, y)

newtype Pack p => UPArray i p = UPArray { unUPArray :: UArray (i, Int) (Item p) }

instance (Pack p, IArray UArray (Item p)) => IArray UPArray p where
    bounds = (fst *** fst) . bounds . unUPArray

    numElements a = numElements (unUPArray a) `div` w where
        w = packSize (unsafeAt a undefined)

    unsafeArray ix es = a where
        w = packSize (unsafeAt a undefined)
        ix' = ((fst ix, 0), (snd ix, w - 1))
        a = UPArray (unsafeArray ix' es')
        es' = do
            (i0, x0) <- es
            zip [i0 * w ..] (itemsUnpack x0)

    unsafeAt a = at where
        w = packSize (at undefined)
        at0 = unsafeAt (unUPArray a)
        at n = packItems [at0 n0 | let base = w*n, n0 <- [base .. (base + w - 1)]]

instance (Pack p, Show (Item p), Show i, Ix i, IArray UArray (Item p)) => Show (UPArray i p) where
    showsPrec p = showsPrec p . unUPArray

test = listArray (0, 2) [('x','y'), ('y', 'z'), ('z', 'x')] :: UPArray Int (Char, Char)
Yeah, that Show instance is fake. Better take a closer look at those w with undefined. I know there is scope types, but I don't like them.