Reasonably painless bidirectional serialization of sums of products

Prickler.hs

{-# LANGUAGE TemplateHaskell, ExistentialQuantification, RankNTypes, GADTs, ImplicitParams, TypeFamilies, DataKinds, TypeOperators, DeriveGeneric, PolyKinds #-}
module Prickler where

{-
Main goals: 
# Minimal indirection, so I didn't want big n-tuples of stuff getting curried into constructors or shit like that
# Magical constructor/eliminator pairing, so I don't have to write ugly unsafe pattern matches or constructor -> Int mappings (unlike the alt combinator in the pickler paper)
# Minimize the amount of metadata traversal that happens during serialization (not 100% there yet)
# Not be ugly to use
-}

import Control.Applicative
import Control.Arrow

import Data.Maybe
import Data.Monoid hiding (Sum, Product, All)
import Data.Binary.Get hiding (Done)
import Data.Binary.Builder
import qualified Data.IntMap as IM
import Data.ByteString.Base16
import qualified Data.ByteString.Lazy as BL

import Data.Int
import Data.Word

newtype K a b = K { unK :: a }
data (:*:) f g x = f x :*: g x
data Exists f = forall a. Exists { getValue :: f a }

newtype Flip b a = Flip { runFlip :: a -> b }
newtype NFlip a ts = NFlip { runNFlip :: ts @-> a }

type family (@->) (ts :: [*]) (r :: *) :: *
type instance '[] @-> r = r
type instance (t ': ts) @-> r = t -> (ts @-> r)

type family Eliminator (css :: [[*]]) (r :: *) :: *
type instance Eliminator '[] r = r
type instance Eliminator (cs ': css) r = cs @-> r -> Eliminator css r



infixr 1 :*

data All f (ts :: [a]) where
  Nil  :: All f '[]
  (:*) :: f x -> All f xs -> All f (x ': xs)

mapAll :: (forall a. f a -> g a) -> All f ts -> All g ts
mapAll f Nil = Nil
mapAll f (a :* as) = f a :* mapAll f as

mapAllF :: (forall a. f a -> b) -> All f ts -> [b]
mapAllF f Nil = []
mapAllF f (a :* as) = f a : mapAllF f as

foldlAll :: (a -> b -> a) -> a -> All (Flip b) ts -> (ts @-> a)
foldlAll f z Nil = z
foldlAll f z (Flip g :* gs) = \x -> foldlAll f (f z (g x)) gs

data Case f i a ts = Case { tag :: i, shape :: All f ts, cons :: ts @-> a } -- does the tag belong here? probably not but I'll leave it there for now
data Data f i a = forall ts. Data { sum :: All (Case f i a) ts, elim :: forall r. a -> EliminatorWrapper ts r }

newtype FunctionWrapper ts r = FunctionWrapper { getFunction :: ts @-> r }
newtype EliminatorWrapper ts r = EliminatorWrapper { getEliminator :: Eliminator ts r }


-- All glory to glguy
apN :: Applicative f => f (ts @-> r) -> All f ts -> f r
apN f Nil = f
apN f (x :* xs) = apN (f <*> x) xs

liftAn :: Applicative f => (ts @-> r) -> All f ts -> f r
liftAn = apN . pure


merge :: All (NFlip r) ts -> Eliminator ts r -> r
merge Nil acc = acc
merge (NFlip x :* xs) acc = merge xs (acc x)


eliminate :: Monoid m => All (K m :*: All (Flip m)) ts -> EliminatorWrapper ts m -> m
eliminate xs = merge (mapAll (\(K i :*: shape) -> NFlip (foldlAll (<>) i shape)) xs) . getEliminator

















data Prickler a = Prickler { get :: Get a, put :: a -> Builder } -- laws: get . put . get == get, put . get . put == put

word8 :: Prickler Word8
word8 = Prickler getWord8 singleton

bool :: Prickler Bool
bool = Prickler (toEnum . fromIntegral <$> getWord8) (singleton . fromIntegral . fromEnum)

char :: Prickler Char
char = Prickler (toEnum . fromIntegral <$> getWord32be) (putWord32be . fromIntegral . fromEnum)


taggedData :: Integral i => Prickler i -> Data Prickler i a -> Prickler a
taggedData (Prickler gi pi) (Data sum elim) = Prickler getter (eliminate (mapAll (adjust pi) sum) . elim)
  where
  adjust :: (i -> Builder) -> Case Prickler i a ts -> (K Builder :*: All (Flip Builder)) ts
  adjust pi (Case i shape _) = K (pi i) :*: mapAll (Flip . put) shape

  ps = IM.fromList (mapAllF (\x -> (fromIntegral (tag x), Exists x)) sum)
  getter = do
    tag <- gi
    case IM.lookup (fromIntegral tag) ps of
      Nothing                           -> fail $ "Invalid tag: " ++ show (fromIntegral tag :: Integer)
      Just (Exists (Case i shape cons)) -> liftAn cons (mapAll get shape)





data Foo 
  = A Word8 Bool
  | B Bool
  | C Char
  | D Word8 Char
  | E
  deriving (Show)

-- Now I just need some automagical code for deriving this thing...
elimFoo :: Foo -> (Word8 -> Bool -> r) -> (Bool -> r) -> (Char -> r) -> (Word8 -> Char -> r) -> r -> r
elimFoo (A w8 x) a b c d e = a w8 x
elimFoo (B x   ) a b c d e = b x
elimFoo (C x   ) a b c d e = c x
elimFoo (D w8 x) a b c d e = d w8 x
elimFoo  E       a b c d e = e


sumFoo
  =  Case 0 (word8 :* bool :* Nil) A
  :* Case 1 (bool  :* Nil) B
  :* Case 2 (char  :* Nil) C
  :* Case 3 (word8 :* char :* Nil) D
  :* Case 4 Nil E
  :* Nil

dataFoo = Data sumFoo (EliminatorWrapper . elimFoo)

foo :: Prickler Foo
foo = taggedData word8 dataFoo

I suppose it might be cleaner to invert the arguments to Case, so the constructor name appears before its arguments...

neat!