bin.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

module Bin (bin) where

import Control.Monad (void)
import Data.Bits (shiftL, (.|.))
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.Char (chr, ord)
import Data.Void (Void)
import Data.Word
import Language.HaskelgetAuthenticationDetailsl.TH
import Language.Haskell.TH.Quote
import Language.Haskell.TH.Syntax
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L

testBytes :: ByteString -> ByteString -> Maybe ByteString
testBytes = ByteString.stripPrefix

fetchBytes :: Int -> ByteString -> Maybe (ByteString, ByteString)
fetchBytes size bs =
  let (before, after) = ByteString.splitAt size bs
   in if ByteString.length before == size then Just (before, after) else Nothing

fetchDecoded :: Decode a => Int -> ByteString -> Maybe (a, ByteString)
fetchDecoded size bs = do
  (a, rest) <- fetchBytes size bs
  val <- decode size (ByteString.unpack a)
  pure (val, rest)

class Decode a where
  decode :: Int -> [Word8] -> Maybe a

instance Decode Char where
  decode n xs | n > 0, all (== 0) (init xs) = Just (chr (fromIntegral (last xs)))
  decode _ _ = Nothing

instance Decode Int where
  decode _ = Just . foldl (\a b -> a `shiftL` 8 .|. fromIntegral b) 0

instance Decode ByteString where
  decode _ = Just . ByteString.pack

data PatElem
  = Char Char
  | Num Integer
  | Binding String
  deriving stock (Show)

data SizeExpr
  = Lit Integer
  | Var String
  | BinOp SizeExpr BinOp SizeExpr
  deriving stock (Show)

data BinOp = Minus | Plus
  deriving stock (Show)

type Pattern = [(PatElem, SizeExpr)]

type Parser = Parsec Void String

sc :: Parser ()
sc = L.space space1 empty empty

lexeme :: Parser a -> Parser a
lexeme = L.lexeme sc

symbol :: String -> Parser ()
symbol = void . L.symbol sc

ident :: Parser String
ident = lexeme ((:) <$> letterChar <*> many alphaNumChar <?> "variable")

patElem :: Parser PatElem
patElem =
  choice
    [ Char <$> between (char '\'') (char '\'') L.charLiteral,
      Num <$> lexeme L.decimal,
      Binding <$> ident
    ]

binOp :: Parser BinOp
binOp =
  choice
    [ Plus <$ symbol "+",
      Minus <$ symbol "-"
    ]

sizeExpr :: Parser SizeExpr
sizeExpr =
  choice
    [ Lit <$> lexeme L.decimal,
      Var <$> ident,
      between (symbol "(") (symbol ")") (BinOp <$> sizeExpr <*> binOp <*> sizeExpr)
    ]

sizedPatElem :: Parser (PatElem, SizeExpr)
sizedPatElem = do
  pe <- patElem
  size <- option (Lit 1) do symbol ":"; sizeExpr
  pure (pe, size)

pattern :: Parser Pattern
pattern = sizedPatElem `sepBy` symbol ","

sizeToExpr :: SizeExpr -> ExpQ
sizeToExpr (Lit n) = litE (IntegerL n)
sizeToExpr (Var v) = varE (mkName v)
sizeToExpr (BinOp a op b) = uInfixE (sizeToExpr a) (opToExpr op) (sizeToExpr b)
  where
    opToExpr Minus = varE (mkName "-")
    opToExpr Plus = varE (mkName "+")

binPat :: String -> Q Pat
binPat src =
  case parse (sc *> pattern <* eof) "" src of
    Left err -> fail (errorBundlePretty err)
    Right elems -> go elems
  where
    go :: [(PatElem, SizeExpr)] -> Q Pat
    go ((Char c, size) : ps) =
      viewP (appE (varE 'fetchDecoded) (sizeToExpr size)) [p| Just ($(litP (CharL c)), $(go ps)) |]
    go ((Num n, size) : ps) =
      viewP (appE (varE 'fetchDecoded) (sizeToExpr size)) [p| Just ($(litP (IntegerL n)) :: Int, $(go ps)) |]
    go ((Binding v, size) : ps) =
      viewP (appE (varE 'fetchDecoded) (sizeToExpr size)) [p| Just ($(varP (mkName v)), $(go ps)) |]
    go [] = viewP (varE 'ByteString.null) [p| True |]

bin = QuasiQuoter {quotePat = binPat}

q = let Just s = parseMaybe sizeExpr "(len-(5+1))" in sizeToExpr s

main = do
  parseTest pattern "'Q', len:4, query:(len-5), 0"