variable-length message binary protocol parser combinators

Parser.hs

module Parser where

import qualified Data.ByteString as B
import Data.Word
import Control.Monad (replicateM)
import System.IO

data Parser s f a
  = Blind  !Int (s -> a)
  | Decide !Int (s -> Parser s f a)
  | Lift (f (Parser s f a))

class Split a where
  split :: Int -> a -> (a,a)

instance Split [a]          where split = splitAt
instance Split B.ByteString where split = B.splitAt

instance Functor f => Functor (Parser s f) where
  fmap f (Blind  i k) = Blind  i (f . k)
  fmap f (Decide i k) = Decide i (fmap f . k)
  fmap f (Lift x    ) = Lift     (fmap (fmap f) x)

instance (Functor f, Split s) => Applicative (Parser s f) where
  pure = Blind 0 . const

  Lift x <*> p = Lift (fmap (<*> p) x)

  Decide i k <*> p =
    Decide i $ \s -> k s <*> p

  Blind i f <*> Decide j g =
    Decide (i+j) $ \s ->
      case split i s of
        (si,sj) -> f si <$> g sj

  Blind i f <*> Blind j g =
    Blind (i+j) $ \s ->
      case split i s of
        (si,sj) -> f si (g sj)

(>>-) :: Functor f => Parser s f a -> (a -> Parser s f b) -> Parser s f b
Lift x     >>- f = Lift (fmap (>>- f) x)
Decide i k >>- f = Decide i ((>>- f) . k)
Blind  i k >>- f = Decide i (f . k)

------------------------------------------------------------------------
-- Example use case ----------------------------------------------------
------------------------------------------------------------------------

word8 :: Parser [Word8] f Word8
word8 = Blind 1 head

word16 :: Parser [Word8] f Word16
word16 = Blind 2 $ \[hi,lo] -> fromIntegral hi * 0x100 + fromIntegral lo

data Message = Message
  { version :: Word8
  , command :: Word8
  , address :: Address
  , port    :: Word16
  }
  deriving Show

data Address
  = IPv4 Word8 Word8 Word8 Word8
  | IPv6 Word16 Word16 Word16 Word16 Word16 Word16 Word16 Word16
  | DomainName [Word8]
  deriving Show

parseAddress :: Parser [Word8] Maybe Address
parseAddress =
  word8 >>- \tag ->
  case tag of
    1 -> IPv4 <$> word8 <*> word8 <*> word8 <*> word8
    3 -> IPv6 <$> word16 <*> word16 <*> word16 <*> word16
              <*> word16 <*> word16 <*> word16 <*> word16
    4 -> word8 >>- \len ->
         DomainName <$> replicateM (fromIntegral len) word8
    _ -> Lift Nothing

parseMessage :: Parser [Word8] Maybe Message
parseMessage = Message <$> word8 <*> word8 <*> parseAddress <*> word16

prompt :: Int -> IO [Word8]
prompt i =
  do putStrLn ("Input list of " ++ show i ++ " Word8")
     readLn

driver :: Show a => Parser [Word8] Maybe a -> IO ()

driver (Blind i k) =
  do xs <- prompt i
     print (show (k xs))

driver (Lift Nothing) = putStrLn "Parse failed"
driver (Lift (Just p)) = driver p

driver (Decide i k) =
  do xs <- prompt i
     driver (k xs)


main =
  do hSetBuffering stdin LineBuffering
     driver parseMessage

{-
*Parser> main
Input list of 3 Word8
[5,1,1]
Input list of 6 Word8
[127,0,0,1,80,80]
"Message {version = 5, command = 1, address = IPv4 127 0 0 1, port = 20560}"

*Parser> main
Input list of 3 Word8
[5,1,4]
Input list of 1 Word8
[7]
Input list of 9 Word8
[1,2,3,4,5,6,7,8,9]
"Message {version = 5, command = 1, address = DomainName [1,2,3,4,5,6,7], port = 2057}"

*Parser> main
Input list of 3 Word8
[10,20,30]
Parse failed

-}