Unboxed polymorphic types

Unboxed.hs

{-# LANGUAGE MagicHash, TypeFamilies #-}

module Data.Unboxed where

import           GHC.Exts

--------------------------------------------------------------------------------
-- Hash (#) types are second-class citizens in GHC (in our case, we can't use
-- them in associated types / type families), so here we add dummy types that
-- stand for "unboxed" variants of some standard GHC types.

data UbxInt    = UbxInt     Int#
data UbxWord   = UbxWord    Word#
data UbxFloat  = UbxFloat   Float#
data UbxDouble = UbxDouble  Double#
data UbxChar   = UbxChar    Char#
data UbxPair a = UbxPair    a a

-- In the rest of the code we just assume these are magically unboxed types.

--------------------------------------------------------------------------------

class Unbox (a :: *) where
    type Unboxed a :: *
    unbox :: a -> Unboxed a
    box   :: Unboxed a -> a

instance Unbox Int where
    type Unboxed Int = UbxInt
    unbox (I# x) = (UbxInt x)
    box (UbxInt i) = I# i

instance Unbox Word where
    type Unboxed Word = UbxWord
    unbox (W# w) = (UbxWord w)
    box (UbxWord w) = W# w

instance Unbox Float where
    type Unboxed Float = UbxFloat
    unbox (F# f) = UbxFloat f
    box (UbxFloat f) = F# f

instance Unbox Double where
    type Unboxed Double = UbxDouble
    unbox (D# x) = UbxDouble x
    box (UbxDouble d) = D# d

instance Unbox Char where
    type Unboxed Char = UbxChar
    unbox (C# c) = UbxChar c
    box (UbxChar c) = C# c

--------------------------------------------------------------------------------

data Point a = Point a a

instance Unbox a => Unbox (Point a) where
  type Unboxed (Point a) = UbxPair (Unboxed a)
  unbox (Point a b) = UbxPair (unbox a) (unbox b)
  box (UbxPair a b) = Point (box a) (box b)

--------------------------------------------------------------------------------

type Point# a = Unboxed (Point a)

distance :: (Unbox a, Floating a) => Point# a -> Point# a -> a
distance p0 p1 =
    let Point x0 y0 = box p0
        Point x1 y1 = box p1
        dx = x1 - x0
        dy = y1 - y0
     in
        sqrt (dx * dx + dy * dy)

-- Now, this by itself doesn't do the trick. We need to remove `unbox` calls in
-- the `distance` function, and for that we need to resolve `Unbox` dictionary
-- in compile time and inline calls, simplify etc.

{-# SPECIALIZE distance :: Point# Float  -> Point# Float  -> Float #-}
{-# SPECIALIZE distance :: Point# Double -> Point# Double -> Double #-}

-- Problem: We should be using unboxed types here. Ideally I should write this
-- instead:
--
-- test1 = distance (Point (0.0 :: Float) 0.0)
--                  (Point 100.0          100.0)
--
test1 :: Float
test1 = distance (UbxPair (UbxFloat 0.0#)   (UbxFloat 0.0#))
                 (UbxPair (UbxFloat 100.0#) (UbxFloat 100.0#))

-- This function is not great for benchmarking this approach though.. It's so
-- simple GHC evaluates the whole thing in compile time:
--
-- (this is -ddump-simpl)
--
-- test1 :: Float
-- test1 = GHC.Types.F# (sqrtFloat# (__float 20000.0))
--
-- So here I'm adding NOINLINE variant to avoid compile time evaluation. Just to
-- see how good the function itself is optimized.
--
{-# NOINLINE distance_noinline #-}
distance_noinline :: (Unbox a, Floating a) => Point# a -> Point# a -> a
distance_noinline p0 p1 =
    let Point x0 y0 = box p0
        Point x1 y1 = box p1
        dx = x1 - x0
        dy = y1 - y0
     in
        sqrt (dx * dx + dy * dy)

{-# SPECIALIZE distance_noinline :: Point# Float  -> Point# Float  -> Float #-}
{-# SPECIALIZE distance_noinline :: Point# Double -> Point# Double -> Double #-}

test2 :: Float
test2 = distance_noinline
                 (UbxPair (UbxFloat 0.0#)   (UbxFloat 0.0#))
                 (UbxPair (UbxFloat 100.0#) (UbxFloat 100.0#))

-- So this sucks. I forgot that GHC doesn't SPECIALIZE without INLINE. I don't
-- understand why.
--
--   Unboxed.hs:115:1: warning:
--       Ignoring useless SPECIALISE pragma for NOINLINE function: ‘distance_noinline’
--
--   Unboxed.hs:116:1: warning:
--       Ignoring useless SPECIALISE pragma for NOINLINE function: ‘distance_noinline’
--
-- So this function is not optimized at all:
--
--   test2 :: Float
--   test2 =
--     distance_noinline
--       @ Float
--       Data.Unboxed.$fUnboxFloat
--       GHC.Float.$fFloatingFloat
--       (Data.Unboxed.test6 `cast` ...)
--       (Data.Unboxed.test4 `cast` ...)
--
-- Here's how it might work in the ideal case: A `distance` specialized on the
-- given types would be created. Then, in the functions, `box` calls would be
-- reduced in compile time as an optimization, so the function body would work
-- on unboxed type.
--
-- This sounds quite similar to worker/wrapper kind of transformations though.
-- Only difference is here we get to say how to unbox and box + what types to
-- specialize on. This might potentially be more reliable.