{-# LANGUAGE TemplateHaskellQuotes #-}
-- |
-- Module:     Data.Bytes.QuasiQuoter
-- Copyright:  (c) 2024 0y2k
-- License:    Apache-2.0
-- Maintainer: 0y2k <0x0y2k@gmail.com>
--
-- Offers QuasiQuoters for byte sequence.

module Data.Bytes.QuasiQuoter
  ( FromBytes(..)
  , fromDigitWithBase
  , strWithBase
  , bin
  , hex
  ) where

import Data.Bits
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import Data.Char
import Data.Maybe (mapMaybe)
import Language.Haskell.TH.Quote

-- $setup
-- >>> :seti -XQuasiQuotes

-- | Offers @StrictByteString -> a@.
class FromBytes a where
  fromBytes :: BS.ByteString -> a

instance FromBytes BS.ByteString where
  fromBytes :: ByteString -> ByteString
fromBytes = ByteString -> ByteString
forall a. a -> a
id
instance FromBytes BL.ByteString where
  fromBytes :: ByteString -> ByteString
fromBytes = ByteString -> ByteString
BL.fromStrict

ord_0, ord_A, ord_a :: Int
ord_0 :: Int
ord_0 = Char -> Int
ord Char
'0'
ord_A :: Int
ord_A = Char -> Int
ord Char
'A'
ord_a :: Int
ord_a = Char -> Int
ord Char
'a'

-- | Accepts @\'0\'..\'9\', \'A\'..\'Z\', \'a\'..\'z\'@.
-- @1 <= base <= 36@
fromDigitWithBase :: Int -> Char -> Maybe Int
fromDigitWithBase :: Int -> Char -> Maybe Int
fromDigitWithBase Int
base
  | Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base Bool -> Bool -> Bool
&& Int
base Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
10 = \Char
c -> let i :: Int
i = Char -> Int
ord Char
c in
    if Int
ord_0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base
      then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_0
      else Maybe Int
forall a. Maybe a
Nothing
  | Int
10 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base Bool -> Bool -> Bool
&& Int
base Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
36 = \Char
c -> let i :: Int
i = Char -> Int
ord Char
c in
    if Int
ord_0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base
      then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_0
      else if Int
ord_A Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_A Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
10
        then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_A Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10
        else if Int
ord_a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
10
          then Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ord_a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10
          else Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Maybe Int -> Char -> Maybe Int
forall a b. a -> b -> a
const Maybe Int
forall a. Maybe a
Nothing

strWithBase' :: Int -> String -> BS.ByteString
strWithBase' :: Int -> String -> ByteString
strWithBase' Int
b
  | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
5 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
b = String -> String -> ByteString
forall a. HasCallStack => String -> a
error String
"strWithBase: illegal b"
  | Bool
otherwise = [Word8] -> ByteString
BS.pack ([Word8] -> ByteString)
-> (String -> [Word8]) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> [Word8]
g ([Bool] -> [Word8]) -> (String -> [Bool]) -> String -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Bool]
forall {a}. Bits a => [a] -> [Bool]
f ([Int] -> [Bool]) -> (String -> [Int]) -> String -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Maybe Int) -> String -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Int -> Char -> Maybe Int
fromDigitWithBase Int
base)
 where
  base :: Int
base = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
b
  f :: [a] -> [Bool]
f [] = []
  f (a
i:[a]
is) = ((Int -> Bool) -> [Int] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
s -> a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
i Int
s) ([Int] -> [Bool]) -> [Int] -> [Bool]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0..Int -> Int
forall a. Enum a => a -> a
pred Int
b]) [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ [a] -> [Bool]
f [a]
is
  g :: [Bool] -> [Word8]
g [] = []
  g [Bool
x0] = [Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7]
  g [Bool
x0,Bool
x1]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6]
  g [Bool
x0,Bool
x1,Bool
x2]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5]
  g [Bool
x0,Bool
x1,Bool
x2,Bool
x3]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x3 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4]
  g [Bool
x0,Bool
x1,Bool
x2,Bool
x3,Bool
x4]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x3 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x4 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3]
  g [Bool
x0,Bool
x1,Bool
x2,Bool
x3,Bool
x4,Bool
x5]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x3 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x4 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x5 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
2]
  g [Bool
x0,Bool
x1,Bool
x2,Bool
x3,Bool
x4,Bool
x5,Bool
x6]
    = [ Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x3 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x4 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x5 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
2
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x6 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
1]
  g (Bool
x0:Bool
x1:Bool
x2:Bool
x3:Bool
x4:Bool
x5:Bool
x6:Bool
x7:[Bool]
xs)
    = ( Bool -> Word8
h Bool
x0 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
7
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
6
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x2 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
5
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x3 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x4 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
3
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x5 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
2
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x6 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
1
    Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Bool -> Word8
h Bool
x7) Word8 -> [Word8] -> [Word8]
forall a. a -> [a] -> [a]
: [Bool] -> [Word8]
g [Bool]
xs
  h :: Bool -> Word8
h = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> (Bool -> Int) -> Bool -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Int
forall a. Enum a => a -> Int
fromEnum

-- | Internal function for quasiquotes. @1 <= b <= 5@
strWithBase :: FromBytes a => Int -> String -> a
strWithBase :: forall a. FromBytes a => Int -> String -> a
strWithBase Int
b = ByteString -> a
forall a. FromBytes a => ByteString -> a
fromBytes (ByteString -> a) -> (String -> ByteString) -> String -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String -> ByteString
strWithBase' Int
b

qqWithBase :: Int -> QuasiQuoter
qqWithBase :: Int -> QuasiQuoter
qqWithBase Int
b = (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
  (\String
s -> [e|strWithBase b s|])
  (String -> String -> Q Pat
forall a. HasCallStack => String -> a
error String
"qq-bytes: cannot use this quasiquoter as a pattern")
  (String -> String -> Q Type
forall a. HasCallStack => String -> a
error String
"qq-bytes: cannot use this quasiquoter as a type")
  (String -> String -> Q [Dec]
forall a. HasCallStack => String -> a
error String
"qq-bytes: cannot use this quasiquoter as a dec")

-- | Binary string into byte sequence.
-- Skips non-binary character.
-- Internally means @'strWithBase' 1@.
--
-- >>> [bin|00100000|] :: BS.ByteString
-- " "
bin :: QuasiQuoter
bin :: QuasiQuoter
bin = Int -> QuasiQuoter
qqWithBase Int
1

-- | Hex string into byte sequence.
-- Skips non-hex character.
-- Internally means @'strWithBase' 4@.
--
-- >>> [hex|48 65 6c 6c 6f|] :: BS.ByteString
-- "Hello"
hex :: QuasiQuoter
hex :: QuasiQuoter
hex = Int -> QuasiQuoter
qqWithBase Int
4