capnp-0.17.0.0: tests/Util.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TypeFamilies #-}

module Util
  ( MsgMetaData (..),
    capnpEncode,
    capnpDecode,
    capnpCompile,
    capnpCanonicalize,
    decodeValue,
    encodeValue,
    aircraftSchemaSrc,
    schemaSchemaSrc,
  )
where

import qualified Capnp.Message as M
import Control.Monad.Trans (lift)
import Control.Monad.Trans.Resource (ResourceT, allocate, runResourceT)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Lazy.Char8 as LBSC8
import System.Directory (removeFile)
import System.Exit (ExitCode (..))
import System.IO
import System.Process hiding (readCreateProcessWithExitCode)
import System.Process.ByteString.Lazy (readCreateProcessWithExitCode)
import Text.Heredoc (there)

aircraftSchemaSrc, schemaSchemaSrc :: String
aircraftSchemaSrc = [there|tests/data/aircraft.capnp|]
schemaSchemaSrc = [there|tests/data/schema.capnp|]

-- | Information about the contents of a capnp message. This is enough
-- to encode/decode both textual and binary forms.
data MsgMetaData = MsgMetaData
  { -- | The source of the schema
    msgSchema :: String,
    -- | The name of the root struct's type
    msgType :: String
  }
  deriving (Show)

capnpCanonicalize :: LBS.ByteString -> IO LBS.ByteString
capnpCanonicalize stdInBytes = do
  (exitStatus, stdOut, stdErr) <-
    readCreateProcessWithExitCode
      (proc "capnp" ["convert", "binary:canonical"])
      stdInBytes
  case exitStatus of
    ExitSuccess -> pure stdOut
    ExitFailure code ->
      fail $
        concat
          [ "capnp convert binary:canonical failed with exit code ",
            show code,
            ":\n",
            show stdErr
          ]

-- | @capnpEncode msg meta@ runs @capnp encode@ on the message, providing
-- the needed metadata and returning the output
capnpEncode :: String -> MsgMetaData -> IO BS.ByteString
capnpEncode msgValue meta = do
  (exitStatus, stdOut, stdErr) <-
    runResourceT $
      interactCapnpWithSchema "encode" (msgSchema meta) (LBSC8.pack msgValue) [msgType meta]
  case exitStatus of
    ExitSuccess -> return (LBS.toStrict stdOut)
    ExitFailure code -> fail ("`capnp encode` failed with exit code " ++ show code ++ ":\n" ++ show stdErr)

-- | @capnpDecode msg meta@ runs @capnp decode@ on the message, providing
-- the needed metadata and returning the output
capnpDecode :: BS.ByteString -> MsgMetaData -> IO String
capnpDecode encodedMsg meta = do
  (exitStatus, stdOut, stdErr) <-
    runResourceT $
      interactCapnpWithSchema "decode" (msgSchema meta) (LBS.fromStrict encodedMsg) [msgType meta]
  case exitStatus of
    ExitSuccess -> return (LBSC8.unpack stdOut)
    ExitFailure code -> fail ("`capnp decode` failed with exit code " ++ show code ++ ":\n" ++ show stdErr)

-- | @capnpCompile msg meta@ runs @capnp compile@ on the schema, providing
-- the needed metadata and returning the output
capnpCompile :: String -> String -> IO BS.ByteString
capnpCompile msgSchema outputArg = do
  (exitStatus, stdOut, stdErr) <-
    runResourceT $
      interactCapnpWithSchema "compile" msgSchema LBSC8.empty ["-o", outputArg]
  case exitStatus of
    ExitSuccess -> return (LBS.toStrict stdOut)
    ExitFailure code -> fail ("`capnp compile` failed with exit code " ++ show code ++ ":\n" ++ show stdErr)

-- | A helper for @capnpEncode@ and @capnpDecode@. Launches the capnp command
-- with the given subcommand (either "encode" or "decode") and metadata,
-- returning handles to its standard in and standard out. This runs inside
-- ResourceT, and sets the handles up to be closed and the process to be reaped
-- when the ResourceT exits.
interactCapnpWithSchema :: String -> String -> LBS.ByteString -> [String] -> ResourceT IO (ExitCode, LBS.ByteString, LBS.ByteString)
interactCapnpWithSchema subCommand msgSchema stdInBytes args = do
  let writeTempFile = runResourceT $ do
        (_, (path, hndl)) <-
          allocate
            (openTempFile "/tmp" "schema.capnp")
            (\(_, hndl) -> hClose hndl)
        lift $ hPutStr hndl msgSchema
        return path
  schemaFile <- snd <$> allocate writeTempFile removeFile
  lift $ readCreateProcessWithExitCode (proc "capnp" ([subCommand, schemaFile] ++ args)) stdInBytes

-- | @'decodeValue' schema typename message@ decodes the value at the root of
-- the message and converts it to text. This is a thin wrapper around
-- 'capnpDecode'.
decodeValue :: String -> String -> M.Message 'M.Const -> IO String
decodeValue schema typename msg =
  let bytes = M.encode msg
   in capnpDecode
        (LBS.toStrict $ BB.toLazyByteString bytes)
        (MsgMetaData schema typename)

-- | @'encodeValue' schema typename value@ encodes the textual value @value@
-- as a capnp message. This is a thin wrapper around 'capnpEncode'.
encodeValue :: String -> String -> String -> IO (M.Message 'M.Const)
encodeValue schema typename value =
  let meta = MsgMetaData schema typename
   in capnpEncode value meta >>= M.decode