binary-v3

Bin.hs

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

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

import Control.Monad (guard)
import Data.Bits (shiftL, (.|.))
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.Char (chr)
import Data.List (foldl')
import Data.Word (Word8)
import Language.Haskell.Exts.Extension qualified as Exts
import Language.Haskell.Exts.Parser qualified as Exts
import Language.Haskell.Meta.Syntax.Translate qualified as Meta
import Language.Haskell.TH
import Language.Haskell.TH.Quote

parsePat :: String -> Pat
parsePat = Meta.toPat . Exts.fromParseResult . Exts.parsePatWithMode mode
  where
    exts = [Exts.ScopedTypeVariables, Exts.ViewPatterns]
    mode = Exts.defaultParseMode {Exts.extensions = [Exts.EnableExtension e | e <- exts]}

class BinarySegment a where
  decodeBinary :: ByteString -> Maybe a

instance BinarySegment Char where
  decodeBinary (ByteString.unpack -> [c]) = Just (chr (fromIntegral c))
  decodeBinary _ = Nothing

instance BinarySegment Word8 where
  decodeBinary (ByteString.unpack -> [w]) = Just w
  decodeBinary _ = Nothing

instance BinarySegment ByteString where
  decodeBinary = Just

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

data Segment = Segment Pat Size
  deriving stock (Show)

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

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

parseSegments :: String -> [Segment]
parseSegments src =
  let ListP ps = parsePat ("[" <> src <> "]")
   in map (uncurry toSegment) (markWithLast ps)
  where
    toSegment :: Bool -> Pat -> Segment
    toSegment True (TildeP pat) = Segment pat Spread
    toSegment False (TildeP pat) = error "spreading is available only for last segment"
    toSegment _ (ViewP size pat) = Segment pat (Limited size)
    toSegment _ pat = Segment pat (Limited (LitE (IntegerL 1)))

split :: Int -> ByteString -> Maybe (ByteString, ByteString)
split size bs
  | ByteString.length bs >= size = Just (ByteString.splitAt size bs)
  | otherwise = Nothing

fetch :: BinarySegment a => Int -> ByteString -> Maybe (a, ByteString)
fetch size bs = do
  (part, rest) <- split size bs
  value <- decodeBinary part
  pure (value, rest)

skip :: Int -> ByteString -> Maybe ByteString
skip size bs = do (_, rest) <- split size bs; pure rest

binPat :: String -> PatQ
binPat = build . parseSegments
  where
    build = \case
      [] -> [p|(ByteString.null -> True)|]
      [Segment WildP Spread] -> [p|_|]
      [Segment pat Spread] ->
        [p|(decodeBinary -> Just $(pure pat))|]
      Segment WildP (Limited size) : segs ->
        [p|(skip $(pure size) -> Just $(build segs))|]
      Segment pat (Limited size) : segs ->
        [p|(fetch $(pure size) -> Just ($(pure pat), $(build segs)))|]

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

Main.hs

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

import Bin (BinaryExtract (..), BinarySegment (..), bin)
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.Char (ord)
import Data.Word (Word8)

data Message
  = Query ByteString
  deriving stock (Show)

newtype PGString = PGString ByteString
  deriving stock (Show)

instance BinarySegment PGString where
  decodeBinary (ByteString.unsnoc -> Just (bs, 0)) = Just (PGString bs)
  decodeBinary _ = Nothing

decodeMessage :: ByteString -> Maybe Message
decodeMessage = \case
  [bin| 'Q', 4 -> len, len-4 -> PGString query |] -> Just (Query query)
  [bin| 'W', 4 -> _, ~(PGString query) |] -> Just (Query query)
  [bin| 'E', 4 -> len, len-5 -> query, 0 :: Word8 |] -> Just (Query query)
  _ -> Nothing

main = do
  print (decodeMessage (mkMsg 'Q'))
  print (decodeMessage (mkMsg 'W'))
  print (decodeMessage (mkMsg 'E'))
  where
    mkMsg l = ByteString.pack [fromIntegral (ord l), 0, 0, 0, 10, 104, 101, 108, 108, 111, 0]

package.yaml

name: hspg

dependencies:
  - base >= 4.14
  - bytestring
  - haskell-src-meta
  - haskell-src-exts
  - template-haskell

executables:
  hspg-exe:
    main: Main.hs