Number-indexed tuples

Tuples.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FunctionalDependencies #-}

import Prelude
import Data.Proxy
import Data.Kind (Type)
import GHC.TypeNats

-- The point of this sketch is to show how we could get this style of tuple
-- index, in such a way that GHC might be able to recover good runtime
-- performance by inlining.

example_0 :: Int
example_0 = at @1 (tuple (42, True))

-- | First we'll need a typical recursive natural number definition so that we
-- can define instances on it (harder to do with the GHC.TypeNats.Nat kind).
data Peano where
  Z :: Peano
  S :: Peano -> Peano

type family ToPeano (n :: Nat) :: Peano where
  ToPeano 0 = 'Z
  ToPeano n = 'S (ToPeano (n - 1))

-- | We'll define the signature of a tuple as a list of its member types.
--
--   (a, b, c) ~ '[a, b, c]
--
type Sig = [Type]

-- | For each signature, we have an "eliminator": this is what you have to give
-- in order to "use" a tuple. For example, if we have (a, b), then if you give
-- a function (a -> b -> r), you can get an r, for any r.
type family Eliminator (sig :: Sig) (r :: Type) :: Type where
  Eliminator '[] r = r
  Eliminator (a ': sig) r = a -> Eliminator sig r

-- | A generic tuple type, parameterized by a signature, in such a way that
-- (hopefully) GHC will inline it well and we can still get decent performance:
-- it's the CPS style definition of a product of the types in the signature.
--
-- The proxy is here for practical reasons: Eliminator is a type family, so GHC
-- needs some way to know what r is.
newtype Tuple (sig :: Sig) = Tuple
  { getTuple :: forall r. Proxy r -> Eliminator sig r -> r }

-- | Given a value, make an eliminator on some signature. This corresponds to
-- ignoring the rest of the signature. It allows us to construct
--
--   \t _ _ _ _ ... _ _ _ -> t
--
class Ignore (sig :: Sig) where
  ignore :: Proxy sig -> t -> Eliminator sig t

instance Ignore '[] where
  {-# INLINE ignore #-}
  ignore _ t = t

instance Ignore rest => Ignore (a ': rest) where
  {-# INLINE ignore #-}
  ignore _ t = \_ -> ignore (Proxy :: Proxy rest) t

-- | For some t, introduce the eliminator for the signature beginning with that
-- type. This could probably be consolidated with Ignore into one typeclass.
class Introduce (t :: Type) (sig :: Sig) where
  introduce :: Proxy t -> Proxy sig -> Eliminator sig t

instance Introduce () '[] where
  introduce _ _ = ()

instance Ignore rest => Introduce a (a ': rest) where
  {-# INLINE introduce #-}
  introduce _ _ = \t -> ignore (Proxy :: Proxy rest) t

-- | How to deconstruct a tuple using an index number. The functional dependency
-- is intuitive: the index and the signature should determine the type.
class Index (index :: Peano) (sig :: Sig) (ty :: Type) | index sig -> ty where
  index :: Proxy index -> Tuple sig -> ty

instance Index 'Z '[] () where
  {-# INLINE index #-}
  index _ (Tuple k) = k Proxy ()

-- | Indexing 0 at the head of the signature: use 'introduce' to make the
-- eliminator, and apply the tuple continuation to it.
instance Introduce a (a ': rest) => Index 'Z (a ': rest) a where
  {-# INLINE index #-}
  index _ (Tuple k) = k Proxy (introduce (Proxy :: Proxy a) (Proxy :: Proxy (a ': rest)))

-- | Indexing greater than 0: index in the smaller tuple that ignores the first
-- element.
instance (WidenEliminator rest, Index n rest b) => Index ('S n) (a ': rest) b where
  {-# INLINE index #-}
  index _ tuple = index (Proxy :: Proxy n) (shortenTuple tuple)

-- | If we can elminate (a, b), then we can eliminate (x, a, b), by ignoring x.
class WidenEliminator (sig :: Sig) where
  widenEliminator :: Proxy sig -> Proxy r -> Eliminator sig r -> Eliminator (a ': sig) r

instance WidenEliminator '[] where
  {-# INLINE widenEliminator #-}
  widenEliminator _ _ r = \_ -> r

instance WidenEliminator rest => WidenEliminator (a ': rest) where
  {-# INLINE widenEliminator #-}
  widenEliminator _ _ elim = \_ -> elim

-- | A tuple is shortened by applying it to a wider eliminator: we have the
-- continuation which expects (a, b, ...), but the smaller tuple will only have
-- an eliminator for (b, ...), but WidenEliminator shows that the latter
-- determines the former, by ignoring the extra element.
shortenTuple :: forall sig a. WidenEliminator sig => Tuple (a ': sig) -> Tuple sig
shortenTuple (Tuple k) =
  Tuple (\proxyR elim -> k proxyR (widenEliminator (Proxy :: Proxy sig) proxyR elim))
{-# INLINE shortenTuple #-}

-- | How to construct a Tuple from a corresponding built-in tuple type.
class FromTuple (ty :: Type) (sig :: Sig) | ty -> sig where
  tuple :: ty -> Tuple sig

-- Would write or generate instances for tuples up to some arbitrary number.
instance FromTuple (a, b) '[a, b] where
  {-# INLINE tuple #-}
  tuple ~(a, b) = Tuple $ \_ k -> k a b

instance FromTuple (a, b, c) '[a, b, c] where
  {-# INLINE tuple #-}
  tuple ~(a, b, c) = Tuple $ \_ k -> k a b c

-- | Use this with a type application on a number.
at :: forall index ty sig. (Index (ToPeano (index - 1)) sig ty) => Tuple sig -> ty
at = index (Proxy @(ToPeano (index - 1)))
{-# INLINE at #-}

example_1 :: Bool
example_1 = at @2 (tuple (42 :: Int, False))

example_2 :: String
example_2 = at @3 (tuple ("hello", "kind", "world"))