Unboxable sigmas à la (unimplemented) Kmett, based on "unzipping" with a rank datastructure (currently super shitty). Please excuse the crap quality of the code, but I was just banging out ideas.

UnboxableSums.hs

{-# LANGUAGE TypeFamilies, GADTs, ConstraintKinds, RankNTypes, TypeOperators, ScopedTypeVariables, FlexibleInstances, FlexibleContexts, MultiParamTypeClasses #-}
module UnboxableSums where

import Prelude hiding (lookup)
import Data.Word
import Data.Bits
import Data.Maybe
import Data.Function
import Data.Constraint
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Data.Proxy
import GHC.Prim
import Unsafe.Coerce

class RankSelect a where
  type Element a :: *
  lookup :: a -> Int -> Element a
  size   :: a -> Int

  rank   :: a -> Element a -> Int -> Int
  select :: a -> Element a -> Int -> Int

  toList   :: a -> [Element a]
  fromList :: [Element a] -> a

-- Super naive instance. Someone write a wavelet tree parametrized on an underlying instance of RankSelect with Element ~ Bool, e.g., RRR!
instance Eq a => RankSelect [a] where
  type Element [a] = a

  lookup = (!!)
  size   = length

  rank   xs x n = length (filter (== x) (take n xs))
  select xs x n = fst (filter ((== x) . snd) (zip [0..] xs) !! n)

  toList   = id
  fromList = id


data MapF f g x where
  MapF :: f a -> MapF f g (g a)

data Const a b = Const a deriving (Eq)

data Sigma f where
  Sigma :: f a -> a -> Sigma f

data (:==) a b where
  Refl :: a :== a

class Eq1 f where
  (===) :: f a -> f b -> Maybe (a :== b)

instance (Eq1 f, Eq b) => Eq (Sigma (MapF f (Const b))) where
  Sigma (MapF x) (Const a) == Sigma (MapF y) (Const b) = isJust (x === y) && a == b


class Forall f (c :: * -> Constraint) where
  witness :: f a -> proxy c -> Dict (c a)


class Eq1 f => Finite f where
  count     :: proxy f -> Int
  toIndex   :: forall a. f a -> Int
  fromIndex :: forall r. Int -> (forall a. f a -> r) -> r

instance Eq1 f => Eq1 (MapF f g) where
  MapF x === MapF y | Just Refl <- x === y = Just Refl
                    | Nothing <- x === y = Nothing

instance Finite f => Finite (MapF f g) where
  count _ = count (Proxy :: Proxy f)
  toIndex (MapF x) = toIndex x
  fromIndex i f = fromIndex i (f . MapF)

data HArray (f :: * -> *) where
  HArray :: V.Vector Any -> HArray f

get :: Finite f => HArray f -> f a -> a
get (HArray v) i = unsafeCoerce (v V.! toIndex i)

tabulate :: Finite f => (forall a. f a -> a) -> HArray f
tabulate f = let r = HArray (V.generate (count r) (\i -> fromIndex i (unsafeCoerce . f))) in r





data SigmaArray f = SigmaArray { components :: HArray (MapF f U.Vector), selectors :: [Sigma (MapF f (Const ()))] }

-- A fancier filter
choose :: Eq1 f => [Sigma f] -> f a -> [a]
choose [] s = []
choose (Sigma x y:xs) s | Just Refl <- x === s = y : choose xs s
                        | Nothing   <- x === s =     choose xs s


toSigmaArray :: forall f. (Forall f U.Unbox, Finite f) => [Sigma f] -> SigmaArray f
toSigmaArray xs = SigmaArray components selectors
  where
  chooser :: MapF f U.Vector a -> a
  chooser (MapF x) | Dict <- witness x (Proxy :: Proxy U.Unbox) = U.fromList (choose xs x)

  components :: HArray (MapF f U.Vector)
  components = tabulate chooser

  selectors :: [Sigma (MapF f (Const ()))]
  selectors = map (\(Sigma _1 _2) -> Sigma (MapF _1) (Const ())) xs


getComponent :: Finite f => SigmaArray f -> f a -> U.Vector a
getComponent (SigmaArray c _) x = get c (MapF x)

index :: forall f. (Forall f U.Unbox, Finite f) => SigmaArray f -> Int -> Sigma f
index (SigmaArray c s) i = case s !! i of q@(Sigma (MapF x) _) | Dict <- witness x (Proxy :: Proxy U.Unbox) -> Sigma x (get c (MapF x) U.! rank s q i)



data AorB a b c where
  A :: AorB a b a
  B :: AorB a b b

left :: a -> Sigma (AorB a b)
left = Sigma A

right :: b -> Sigma (AorB a b)
right = Sigma B

instance (U.Unbox a, U.Unbox b) => Forall (AorB a b) U.Unbox where
  witness A _ = Dict
  witness B _ = Dict

instance Eq1 (AorB a b) where
  A === A = Just Refl
  B === B = Just Refl
  _ === _ = Nothing

instance Finite (AorB a b) where
  count _ = 2
  
  toIndex A = 0
  toIndex B = 1

  fromIndex 0 f = f A
  fromIndex 1 f = f B

test = toSigmaArray [left (5 :: Int), right True, left 7, left 1, right False]

tester i = case index test i of
  Sigma A x -> show x
  Sigma B y -> show y