binary-v2

Bin.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}

module Bin (BinarySegment (..), bin) where

import Control.Applicative ((<|>))
import Control.Monad (guard)
import Data.Bits
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.Char qualified as Char
import Data.List (foldl')
import Data.Proxy
import Language.Haskell.Meta.Parse
import Language.Haskell.TH
import Language.Haskell.TH.Quote

class BinarySegment a where
  defaultSize :: Proxy a -> Int
  decodeBinary :: Proxy a -> ByteString -> Maybe a
  encodeBinary :: Proxy a -> Maybe Int -> a -> ByteString

instance BinarySegment Char where
  defaultSize _ = 1
  decodeBinary _ bs
    | not (ByteString.null bs),
      let ws = ByteString.unpack bs,
      all (== 0) (init ws) =
      Just (Char.chr (fromIntegral (last ws)))
  decodeBinary _ _ = Nothing
  encodeBinary _ Nothing a = ByteString.singleton (fromIntegral (Char.ord a))
  encodeBinary _ (Just n) a | n < 1 = ByteString.empty
  encodeBinary p (Just n) a = ByteString.replicate (n - 1) 0 <> encodeBinary p Nothing a

instance BinarySegment ByteString where
  defaultSize _ = 1
  decodeBinary _ = Just
  encodeBinary _ Nothing bs = bs
  encodeBinary _ (Just n) bs = ByteString.take n bs

instance BinarySegment Int where
  defaultSize _ = 1
  decodeBinary _ bs =
    let ws = ByteString.unpack bs
     in Just (foldl' (\a b -> a `shiftL` 8 .|. fromIntegral b) 0 ws)

  encodeBinary p n x = go (case n of Nothing -> 1; Just n -> n) x []
    where
      go 0 x r = ByteString.pack r
      go n x r = go (n - 1) (x `shiftR` 8) (fromIntegral (x .&. 255) : r)

data Segment
  = Int Int
  | Char Char
  | String String
  | Variable Name
  deriving stock (Show)

data Size = Limited Exp | Rest deriving stock (Show)

pattern ConsE <- ConE ((== '(:)) -> True)

pattern UnitE <- ConE ((== mkName "()") -> True)

parsePattern :: String -> [(Segment, Size, Maybe Type)]
parsePattern source =
  case parseExp ("[" <> source <> "]") of
    Left err -> error err
    Right (ListE segments) -> either error id (traverse (uncurry translate) (markWithLast segments))
    _ -> error "invalid format"
  where
    translate :: Exp -> Bool -> Either String (Segment, Size, Maybe Type)
    translate e isLast = do
      (e, mty) <- extractType e
      (e, msize) <- extractSize isLast e
      seg <- extractSegment e
      let mty' = mty <|> defaultSegType seg
      let size = case (msize, mty') of
            (Just size, _) -> size
            (Nothing, Just ty) -> Limited (mkDefaultSize ty)
            (Nothing, Nothing) -> Limited (LitE (IntegerL 1))
      pure (seg, size, mty')

    extractType :: Exp -> Either String (Exp, Maybe Type)
    extractType (SigE e ty) = Right (e, Just ty)
    extractType e = Right (e, Nothing)

    extractSize :: Bool -> Exp -> Either String (Exp, Maybe Size)
    extractSize True (UInfixE e ConsE UnitE) = Right (e, Just Rest)
    extractSize False (UInfixE e ConsE UnitE) = Left "Unlimited size can be used only for last segment"
    extractSize _ (UInfixE e ConsE size) = Right (e, Just (Limited size))
    extractSize _ e = Right (e, Nothing)

    extractSegment :: Exp -> Either String Segment
    extractSegment (LitE (IntegerL i)) = Right (Int (fromInteger i))
    extractSegment (LitE (CharL c)) = Right (Char c)
    extractSegment (LitE (StringL s)) = Right (String s)
    extractSegment (VarE var) = Right (Variable var)
    extractSegment e = Left ("Invalid binary pattern segment: " <> show e)

    defaultSegType :: Segment -> Maybe Type
    defaultSegType (Int _) = Just (ConT ''Int)
    defaultSegType (Char _) = Just (ConT ''Char)
    defaultSegType (String _) = Just (ConT ''ByteString)
    defaultSegType (Variable _) = Nothing

    mkDefaultSize :: Type -> Exp
    mkDefaultSize ty = AppE (VarE 'defaultSize) (SigE (ConE 'Proxy) (AppT (ConT ''Proxy) ty))

    markWithLast :: [a] -> [(a, Bool)]
    markWithLast = foldr (\x -> \case [] -> [(x, True)]; xs -> (x, False) : xs) []

fetch :: BinarySegment a => Proxy a -> Int -> ByteString -> Maybe (a, ByteString)
fetch _ size bs = do
  let (part, rest) = ByteString.splitAt size bs
  guard (ByteString.length part == size)
  value <- decodeBinary Proxy part
  pure (value, rest)

mkProxy :: Maybe Type -> ExpQ
mkProxy Nothing = [e|Proxy|]
mkProxy (Just ty) = [e|Proxy :: Proxy $(pure ty)|]

binPat :: String -> PatQ
binPat = build . parsePattern
  where
    build :: [(Segment, Size, Maybe Type)] -> PatQ
    build [] = [p|(ByteString.null -> True)|]
    build [(seg, Rest, mty)] =
      [p|(decodeBinary $(mkProxy mty) -> Just $(segToPat seg))|]
    build ((seg, Limited size, mty) : segs) =
      [p|(fetch $(mkProxy mty) $(pure size) -> Just ($(segToPat seg), $(build segs)))|]

    segToPat :: Segment -> PatQ
    segToPat (Int i) = litP (IntegerL (fromIntegral i))
    segToPat (Char c) = litP (CharL c)
    segToPat (String s) = [p|((== fromString $(litE (StringL s))) -> True)|]
    segToPat (Variable v) = varP v

binExp :: String -> ExpQ
binExp = build . parsePattern
  where
    build [] = [e|ByteString.empty|]
    build ((seg, size, mty) : segs) =
      [e|encodeBinary $(mkProxy mty) $size' $(segToExp seg) <> $(build segs)|]
      where
        size' = case size of Limited s -> [e|Just $(pure s)|]; Rest -> [e|Nothing|]

    segToExp :: Segment -> ExpQ
    segToExp (Int i) = litE (IntegerL (fromIntegral i))
    segToExp (Char c) = litE (CharL c)
    segToExp (String s) = [e|fromString $(litE (StringL s))|]
    segToExp (Variable v) = varE v

bin :: QuasiQuoter
bin =
  QuasiQuoter
    { quotePat = binPat,
      quoteExp = binExp,
      quoteType = error "unimplemented: bin type",
      quoteDec = error "unimplemented: bin declaration"
    }

main = do
  let q = "'Q', _len:4, query:() :: PGString"
  traverse print (parsePattern q)

Main.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ViewPatterns #-}

import Bin (BinarySegment (..), bin)
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.String (fromString)

data Message
  = Query PGString
  deriving stock (Show)

newtype PGString = PGString ByteString
  deriving stock (Show)

pglen :: PGString -> Int
pglen (PGString bs) = ByteString.length bs

instance BinarySegment PGString where
  defaultSize _ = 1
  decodeBinary _ bs = case ByteString.unsnoc bs of
    Just (bs, 0) -> Just (PGString bs)
    _ -> Nothing
  encodeBinary _ _ (PGString bs) = bs <> ByteString.singleton 0

encodeMessage :: Message -> ByteString
encodeMessage (Query query) =
  let size = pglen query + 5
   in [bin| 'Q', size:4, query:() |]

decodeMessage :: ByteString -> Maybe Message
decodeMessage [bin| 'Q', _len:4 :: Int, query:() |] = Just (Query query)
decodeMessage _ = Nothing

main = do
  let query = PGString (fromString "hello")
  let msg = encodeMessage (Query query)
  print (ByteString.unpack msg)
  print (decodeMessage msg)

-- [81,0,0,0,10,104,101,108,108,111,0]
-- Just (Query (PGString "hello"))