An exercise to learn a bit of type level programming in the style of Servant with kinds

ByteStringCalculator.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Service.TCP.Route where

import Control.Monad (guard)
import Data.Attoparsec.ByteString.Char8 qualified as A
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as B
import Data.Data (Proxy (..))
import Data.Kind (Type)
import GHC.TypeLits (KnownNat, natVal, Natural)
import Prelude

-- | Codecs to parse and render arguments and results from and to ByteString
class Codec a where
    parseArg :: ByteString -> Maybe (a, ByteString)
    renderResult :: a -> ByteString

-- | Parse a value of type 'a' from a ByteString, skipping spaces after the value
parseAndSkipSpaces :: A.Parser a -> ByteString -> Maybe (a, ByteString)
parseAndSkipSpaces parser input = case parse input of
    A.Done rest result -> Just (result, rest)
    _ -> Nothing
  where
    parse = A.parse $ do
        r <- parser
        A.skipWhile (== ' ')
        pure r

instance Codec Int where
    parseArg = parseAndSkipSpaces A.decimal

    renderResult = B.pack . show

instance Codec Double where
    parseArg = parseAndSkipSpaces A.double

    renderResult = B.pack . show

instance Codec ByteString where
    parseArg =
        parseAndSkipSpaces
            $ A.takeWhile (\y -> not $ elem y (" \n" :: String))
    renderResult = id

----- Application --------------------------------------------

-- | A kind that represents a route with a number, a list of arguments and a result
data RouteK n as a = RouteT n as a

-- | Specialized version of RouteT with the number of kind Natural
type Route n as a = RouteT (n :: Natural) (as :: [Type]) (a :: Type)

-- | Extract a the type of the function we have to provide
type family Function route where
    Function (Route n '[] r) = r
    Function (Route n (x ': xs) r) = x -> Function (Route n xs r)

-- | Result of evaluating a route, can fail if was not possible to parse one of the arguments
-- The failure contains the index of the argument that failed to parse
data EvalResult = EvalSuccess ByteString | EvalFailure Int

-- | How to evaluate a route given a bytestring input
class Eval n as r where
    evalRoute
        :: Proxy (Route n as r)
        -- ^ the route we are evaluating
        -> Function (Route n as r)
        -- ^ the function to apply to the arguments, the type is computed
        -- by the type family 'Function', from the route
        -> Int
        -- ^ an index to keep track of the argument we are parsing
        -> ByteString
        -- ^ the input to parse
        -> EvalResult
        -- ^ the result of the operation

instance Codec r => Eval n '[] r where
    evalRoute _ r _ _ = EvalSuccess $ renderResult r

instance (Eval n as r, Codec a) => Eval n (a ': as) r where
    evalRoute _ f n input = case parseArg @a input of
        Just (a, rest) -> evalRoute (Proxy @(Route n as r)) (f a) (succ n) rest
        Nothing -> EvalFailure n

-- | A function that tries to parse the route number and then calls 'evalRoute'
-- if the route number matches the one we are handling
tryRoute
    :: forall n as r
     . (KnownNat n, Eval n as r)
    => Proxy (Route n as r)
    -> Function (Route n as r)
    -> ByteString
    -> Maybe Result
tryRoute p f input = case parseArg @Int input of
    Just (m, rest) -> do
        guard $ fromIntegral m == natVal (Proxy @n)
        pure $ case evalRoute p f 1 rest of
            EvalSuccess r -> Success r
            EvalFailure n -> FailedToParseArg n
    Nothing -> Just FailedToParseRoute

-- | A new data kind to hold the routes
data Routes a b = a :<<|> b | EndOfRoutes

infixr 9 :<<|>

-- | Isomorphic to Either, collects the operations when used recursively
data a :<|> b = a :<|> b

-- | The end of the operations
data End = End

infixr 5 :<|>

-- | A type family to compute the functions we have to provide for each route
type family Functions (x :: Routes a b) :: Type where
    Functions (a :<<|> b) = Function a :<|> Functions b
    Functions EndOfRoutes = End

data Result
    = -- | Failed to parse route
      FailedToParseRoute
    | -- | Failed to parse the argument numbered as
      FailedToParseArg Int
    | -- | No route for the given number
      NoRouteMatched
    | -- | The result of the requested operation
      Success ByteString
    deriving (Show, Eq)

{-| An application is a function from ByteString to ByteString that can fail
Nothing signals that the route was not matched
-}
type Application a = ByteString -> a

---------- Server --------------------------------------------



{-| A type class for creating a server from the operations, this will take care
of the routing.
-}
class Server (routes :: Routes a b) where
    server :: Proxy routes -> Functions routes -> Application Result

instance Server EndOfRoutes where
    server _ _ _ = NoRouteMatched

instance
    ( Eval n as r
    , KnownNat n
    , Server restOfOps
    )
    => Server (Route n as r :<<|> restOfOps)
    where
    server _ (op :<|> ops) input =
        case tryRoute (Proxy @(Route n as r)) op input of
            Just r -> r
            Nothing -> server (Proxy @restOfOps) ops input

serve :: Server ops => Proxy ops -> Functions ops -> Application Result
serve = server -- start the routing at 0 index

-- Test -------

-- A bunch of operations we want to support
type Sum = Route 0 '[Int, Int] Int
type Product = Route 1 '[Int, Int] Int
type Division = Route 2 '[Double, Double] Double
type Negate = Route 3 '[Int] Int
type Reverse = Route 4 '[ByteString] ByteString

type API =
    Sum
        :<<|> Product
        :<<|> Division
        :<<|> Negate
        :<<|> Reverse
        :<<|> EndOfRoutes

-- The implementation of the operations, note the signature of the functions
-- that can be computed from the type family 'Function'
sumOp :: Function Sum -- Int -> Int -> Int
sumOp a b = a + b

productOp :: Function Product -- Int -> Int -> Int
productOp a b = a * b

divisionOp :: Function Division -- Double -> Double -> Double
divisionOp a b = a / b

negateOp :: Function Negate -- Int -> Int
negateOp a = -a

reverseOp :: Function Reverse -- ByteString -> ByteString
reverseOp = B.reverse

service :: Functions API
service = sumOp :<|> productOp :<|> divisionOp :<|> negateOp :<|> reverseOp :<|> End

api :: Proxy API
api = Proxy @API

assert :: (Eq b, Applicative f, Show b) => b -> b -> f ()
assert x y = if x == y then pure () else error $ show (x, y)

test :: IO ()
test = do
    assert (serve api service "0 32 10\n") (Success "42")
    assert (serve api service "1 32 10\n") (Success "320")
    assert (serve api service "2 32 10\n") (Success "3.2")
    assert (serve api service "3 32\n") (Success "-32")
    assert (serve api service "4 hello\n") (Success "olleh")
    assert (serve api service "5 hello\n") NoRouteMatched
    assert (serve api service "0 hello\n") (FailedToParseArg 1)
    assert (serve api service "1 32 i\n") (FailedToParseArg 2)
    assert (serve api service "hello\n") FailedToParseRoute