{-# LANGUAGE GADTs #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Test.Cardano.Ledger.Binary.Cuddle.GenerateCBOR (
  generateCBORMain,
) where

import Codec.CBOR.Cuddle.CBOR.Gen
import Codec.CBOR.Cuddle.CBOR.Validator
import Codec.CBOR.Cuddle.CBOR.Validator.Trace
import Codec.CBOR.Cuddle.CDDL
import Codec.CBOR.Cuddle.CDDL.Custom.Generator
import qualified Codec.CBOR.Cuddle.Huddle as Cuddle
import Codec.CBOR.Cuddle.IndexMappable
import qualified Codec.CBOR.Term as CBOR
import qualified Codec.CBOR.Write as CBOR
import Control.Monad (forM_, when)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base16 as Base16
import qualified Data.ByteString.Char8 as BS8
import Data.Maybe (fromMaybe, isJust)
import qualified Data.Text as T
import Options.Applicative ((<**>))
import qualified Options.Applicative as Opt
import System.Exit (die)
import System.IO
import Test.AntiGen
import Test.Cardano.Ledger.Binary.Cuddle
import Test.QuickCheck (generate)
import Test.QuickCheck.Gen (unGen)
import Test.QuickCheck.Random (mkQCGen)

data GenerateCBOROpts = GenerateCBOROpts
  { GenerateCBOROpts -> [Text]
gcboRuleNames :: ![T.Text]
  , GenerateCBOROpts -> Maybe Int
gcboZap :: !(Maybe Int)
  , GenerateCBOROpts -> Int
gcboCount :: !Int
  , GenerateCBOROpts -> Maybe Int
gcboSeed :: !(Maybe Int)
  , GenerateCBOROpts -> Bool
gcboBinary :: !Bool
  }

optsParser :: Opt.Parser GenerateCBOROpts
optsParser :: Parser GenerateCBOROpts
optsParser =
  [Text] -> Maybe Int -> Int -> Maybe Int -> Bool -> GenerateCBOROpts
GenerateCBOROpts
    ([Text]
 -> Maybe Int -> Int -> Maybe Int -> Bool -> GenerateCBOROpts)
-> Parser [Text]
-> Parser
     (Maybe Int -> Int -> Maybe Int -> Bool -> GenerateCBOROpts)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Text -> Parser [Text]
forall a. Parser a -> Parser [a]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
Opt.some
      ( Mod ArgumentFields Text -> Parser Text
forall s. IsString s => Mod ArgumentFields s -> Parser s
Opt.strArgument (Mod ArgumentFields Text -> Parser Text)
-> Mod ArgumentFields Text -> Parser Text
forall a b. (a -> b) -> a -> b
$
          String -> Mod ArgumentFields Text
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
Opt.metavar String
"RULE_NAME..."
            Mod ArgumentFields Text
-> Mod ArgumentFields Text -> Mod ArgumentFields Text
forall a. Semigroup a => a -> a -> a
<> String -> Mod ArgumentFields Text
forall (f :: * -> *) a. String -> Mod f a
Opt.help String
"CDDL rule names to generate CBOR for"
      )
    Parser (Maybe Int -> Int -> Maybe Int -> Bool -> GenerateCBOROpts)
-> Parser (Maybe Int)
-> Parser (Int -> Maybe Int -> Bool -> GenerateCBOROpts)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser Int -> Parser (Maybe Int)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
Opt.optional
      ( ReadM Int -> Mod OptionFields Int -> Parser Int
forall a. ReadM a -> Mod OptionFields a -> Parser a
Opt.option ReadM Int
forall a. Read a => ReadM a
Opt.auto (Mod OptionFields Int -> Parser Int)
-> Mod OptionFields Int -> Parser Int
forall a b. (a -> b) -> a -> b
$
          String -> Mod OptionFields Int
forall (f :: * -> *) a. HasName f => String -> Mod f a
Opt.long String
"zap"
            Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
Opt.metavar String
"N"
            Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. String -> Mod f a
Opt.help String
"Generate corrupted (zapped) CBOR with N mistakes"
      )
    Parser (Int -> Maybe Int -> Bool -> GenerateCBOROpts)
-> Parser Int -> Parser (Maybe Int -> Bool -> GenerateCBOROpts)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReadM Int -> Mod OptionFields Int -> Parser Int
forall a. ReadM a -> Mod OptionFields a -> Parser a
Opt.option
      ReadM Int
forall a. Read a => ReadM a
Opt.auto
      ( String -> Mod OptionFields Int
forall (f :: * -> *) a. HasName f => String -> Mod f a
Opt.long String
"count"
          Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> Char -> Mod OptionFields Int
forall (f :: * -> *) a. HasName f => Char -> Mod f a
Opt.short Char
'n'
          Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
Opt.metavar String
"N"
          Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> Int -> Mod OptionFields Int
forall (f :: * -> *) a. HasValue f => a -> Mod f a
Opt.value Int
1
          Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> Mod OptionFields Int
forall a (f :: * -> *). Show a => Mod f a
Opt.showDefault
          Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. String -> Mod f a
Opt.help String
"Number of samples to generate per rule"
      )
    Parser (Maybe Int -> Bool -> GenerateCBOROpts)
-> Parser (Maybe Int) -> Parser (Bool -> GenerateCBOROpts)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser Int -> Parser (Maybe Int)
forall (f :: * -> *) a. Alternative f => f a -> f (Maybe a)
Opt.optional
      ( ReadM Int -> Mod OptionFields Int -> Parser Int
forall a. ReadM a -> Mod OptionFields a -> Parser a
Opt.option ReadM Int
forall a. Read a => ReadM a
Opt.auto (Mod OptionFields Int -> Parser Int)
-> Mod OptionFields Int -> Parser Int
forall a b. (a -> b) -> a -> b
$
          String -> Mod OptionFields Int
forall (f :: * -> *) a. HasName f => String -> Mod f a
Opt.long String
"seed"
            Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. HasMetavar f => String -> Mod f a
Opt.metavar String
"SEED"
            Mod OptionFields Int
-> Mod OptionFields Int -> Mod OptionFields Int
forall a. Semigroup a => a -> a -> a
<> String -> Mod OptionFields Int
forall (f :: * -> *) a. String -> Mod f a
Opt.help String
"Fixed random seed for reproducibility"
      )
    Parser (Bool -> GenerateCBOROpts)
-> Parser Bool -> Parser GenerateCBOROpts
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Mod FlagFields Bool -> Parser Bool
Opt.switch
      ( String -> Mod FlagFields Bool
forall (f :: * -> *) a. HasName f => String -> Mod f a
Opt.long String
"binary"
          Mod FlagFields Bool -> Mod FlagFields Bool -> Mod FlagFields Bool
forall a. Semigroup a => a -> a -> a
<> String -> Mod FlagFields Bool
forall (f :: * -> *) a. String -> Mod f a
Opt.help String
"Output raw CBOR bytes instead of hex encoding"
      )

generateCBORMain :: Cuddle.Huddle -> IO ()
generateCBORMain :: Huddle -> IO ()
generateCBORMain Huddle
huddle = do
  opts <-
    ParserInfo GenerateCBOROpts -> IO GenerateCBOROpts
forall a. ParserInfo a -> IO a
Opt.execParser (ParserInfo GenerateCBOROpts -> IO GenerateCBOROpts)
-> ParserInfo GenerateCBOROpts -> IO GenerateCBOROpts
forall a b. (a -> b) -> a -> b
$
      Parser GenerateCBOROpts
-> InfoMod GenerateCBOROpts -> ParserInfo GenerateCBOROpts
forall a. Parser a -> InfoMod a -> ParserInfo a
Opt.info
        (Parser GenerateCBOROpts
optsParser Parser GenerateCBOROpts
-> Parser (GenerateCBOROpts -> GenerateCBOROpts)
-> Parser GenerateCBOROpts
forall (f :: * -> *) a b. Applicative f => f a -> f (a -> b) -> f b
<**> Parser (GenerateCBOROpts -> GenerateCBOROpts)
forall a. Parser (a -> a)
Opt.helper)
        ( InfoMod GenerateCBOROpts
forall a. InfoMod a
Opt.fullDesc
            InfoMod GenerateCBOROpts
-> InfoMod GenerateCBOROpts -> InfoMod GenerateCBOROpts
forall a. Semigroup a => a -> a -> a
<> String -> InfoMod GenerateCBOROpts
forall a. String -> InfoMod a
Opt.progDesc String
"Generate CBOR data from Cardano Ledger CDDL rules"
            InfoMod GenerateCBOROpts
-> InfoMod GenerateCBOROpts -> InfoMod GenerateCBOROpts
forall a. Semigroup a => a -> a -> a
<> String -> InfoMod GenerateCBOROpts
forall a. String -> InfoMod a
Opt.header String
"generate-cbor - CBOR data generator from Cardano Ledger CDDL specifications"
        )
  case resolveHuddle huddle of
    Left String
err -> String -> IO ()
forall a. String -> IO a
die (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Failed to resolve CDDL: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
err
    Right CTreeRoot MonoReferenced
root -> do
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (GenerateCBOROpts -> Bool
gcboBinary GenerateCBOROpts
opts) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> Bool -> IO ()
hSetBinaryMode Handle
stdout Bool
True
      let env :: HuddleEnv
env = HuddleEnv {heTwiddle :: Bool
heTwiddle = Bool
True, heRoot :: CTreeRoot MonoReferenced
heRoot = CTreeRoot MonoReferenced
root}
          multipleRules :: Bool
multipleRules = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (GenerateCBOROpts -> [Text]
gcboRuleNames GenerateCBOROpts
opts) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
      [Text] -> (Text -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (GenerateCBOROpts -> [Text]
gcboRuleNames GenerateCBOROpts
opts) ((Text -> IO ()) -> IO ()) -> (Text -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Text
ruleName -> do
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
multipleRules Bool -> Bool -> Bool
&& Bool -> Bool
not (GenerateCBOROpts -> Bool
gcboBinary GenerateCBOROpts
opts)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
            String
"# " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
ruleName
        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. GenerateCBOROpts -> Int
gcboCount GenerateCBOROpts
opts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ GenerateCBOROpts -> HuddleEnv -> Text -> Int -> IO ()
emitSample GenerateCBOROpts
opts HuddleEnv
env Text
ruleName

emitSample :: GenerateCBOROpts -> HuddleEnv -> T.Text -> Int -> IO ()
emitSample :: GenerateCBOROpts -> HuddleEnv -> Text -> Int -> IO ()
emitSample GenerateCBOROpts
opts HuddleEnv
env Text
ruleName Int
sampleIx = do
  let cborGen :: AntiGen Term
cborGen = GenConfig -> CBORGen Term -> AntiGen Term
forall a. GenConfig -> CBORGen a -> AntiGen a
runCBORGen (HuddleEnv -> GenConfig
toGenConfig HuddleEnv
env) (HasCallStack => Name -> CBORGen Term
Name -> CBORGen Term
generateFromName (Text -> Name
Name Text
ruleName))
      gen :: Gen (ZapResult Term)
gen = Int -> AntiGen Term -> Gen (ZapResult Term)
forall a. Int -> AntiGen a -> Gen (ZapResult a)
zapAntiGenResult (Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (GenerateCBOROpts -> Maybe Int
gcboZap GenerateCBOROpts
opts)) AntiGen Term
cborGen
  result <- case GenerateCBOROpts -> Maybe Int
gcboSeed GenerateCBOROpts
opts of
    Maybe Int
Nothing -> Gen (ZapResult Term) -> IO (ZapResult Term)
forall a. Gen a -> IO a
generate Gen (ZapResult Term)
gen
    Just Int
seed -> ZapResult Term -> IO (ZapResult Term)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ZapResult Term -> IO (ZapResult Term))
-> ZapResult Term -> IO (ZapResult Term)
forall a b. (a -> b) -> a -> b
$ Gen (ZapResult Term) -> QCGen -> Int -> ZapResult Term
forall a. Gen a -> QCGen -> Int -> a
unGen Gen (ZapResult Term)
gen (Int -> QCGen
mkQCGen (Int
seed Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sampleIx)) Int
30
  writeSample opts env ruleName sampleIx result

writeSample :: GenerateCBOROpts -> HuddleEnv -> T.Text -> Int -> ZapResult CBOR.Term -> IO ()
writeSample :: GenerateCBOROpts
-> HuddleEnv -> Text -> Int -> ZapResult Term -> IO ()
writeSample GenerateCBOROpts
opts HuddleEnv
env Text
ruleName Int
sampleIx ZapResult {Term
zrValue :: Term
zrValue :: forall a. ZapResult a -> a
zrValue, Int
zrZapped :: Int
zrZapped :: forall a. ZapResult a -> Int
zrZapped}
  | Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (GenerateCBOROpts -> Maybe Int
gcboZap GenerateCBOROpts
opts) =
      if Int
zrZapped Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then String -> IO ()
warn String
"produced no corruptions"
        else case Either ValidateCBORError (Evidenced ValidationTrace)
validation of
          Left ValidateCBORError
e -> String -> IO ()
warn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"could not run validation because of malformed bytes\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ValidateCBORError -> String
forall a. Show a => a -> String
show ValidateCBORError
e
          Right (Evidenced SValidity v
SValid ValidationTrace v
_) -> String -> IO ()
warn String
"produced a value that is still valid"
          Right (Evidenced SValidity v
SInvalid ValidationTrace v
_) -> IO ()
outputBinary
  | Bool
otherwise = IO ()
outputBinary
  where
    outputBinary :: IO ()
outputBinary
      | GenerateCBOROpts -> Bool
gcboBinary GenerateCBOROpts
opts = Handle -> ByteString -> IO ()
BS.hPut Handle
stdout ByteString
bs
      | Bool
otherwise = ByteString -> IO ()
BS8.putStrLn (ByteString -> ByteString
Base16.encode ByteString
bs)
    bs :: ByteString
bs = Encoding -> ByteString
CBOR.toStrictByteString (Term -> Encoding
CBOR.encodeTerm Term
zrValue)
    validation :: Either ValidateCBORError (Evidenced ValidationTrace)
validation = HasCallStack =>
ByteString
-> Name
-> CTreeRoot ValidatorPhase
-> Either ValidateCBORError (Evidenced ValidationTrace)
ByteString
-> Name
-> CTreeRoot ValidatorPhase
-> Either ValidateCBORError (Evidenced ValidationTrace)
validateCBOR ByteString
bs (Text -> Name
Name Text
ruleName) (CTreeRoot MonoReferenced -> CTreeRoot ValidatorPhase
forall {k} (f :: k -> *) (i :: k) (j :: k).
IndexMappable f i j =>
f i -> f j
mapIndex (HuddleEnv -> CTreeRoot MonoReferenced
heRoot HuddleEnv
env))
    warn :: String -> IO ()
warn String
reason =
      Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
        String
"Warning: zap "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
reason
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" for rule "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
ruleName
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" (sample "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
sampleIx
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
")"