Generator.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Core generator GADT and evaluation machinery.
--
-- Generators produce typed Haskell values by communicating with the Hegel
-- server over CBOR. The 'Generator' type is a GADT with six constructors
-- reflecting different generation strategies:
--
-- * 'Basic' holds a raw CBOR schema and a transform; the server generates the
--   value in one round-trip.
-- * 'Mapped' wraps a source generator and a transform function.
-- * 'FlatMapped' wraps a source and a function returning a generator.
-- * 'Filtered' wraps a source and a predicate (up to 'maxFilterAttempts' retries).
-- * 'CompositeList' uses the collection protocol for non-basic list elements.
-- * 'Composite' wraps a generation function inside a labelled span.
module Hegel.Generator
  ( -- * Generator type
    Generator (..)

    -- * Evaluation
  , draw

    -- * Combinators
  , gmap
  , flatMap
  , gfilter

    -- * Introspection
  , isBasic
  , asBasic

    -- * Span helpers
  , group
  , discardableGroup

    -- * Collection protocol
  , Collection (..)
  , newCollection
  , collectionMore
  , collectionReject

    -- * Span label constants
  , labelList
  , labelListElement
  , labelSet
  , labelSetElement
  , labelMap
  , labelMapEntry
  , labelTuple
  , labelOneOf
  , labelOptional
  , labelFixedDict
  , labelFlatMap
  , labelFilter
  , labelMapped
  , labelSampledFrom
  , labelEnumVariant

    -- * Constants
  , maxFilterAttempts
  ) where

import Codec.CBOR.Term (Term (..))
import Control.Exception (SomeException, catch, throwIO)
import Control.Monad (unless)
import Data.IORef

import Hegel.Client
  ( DataExhausted (..)
  , TestCase (..)
  , assume
  , generateFromSchema
  , startSpan
  , stopSpan
  )
import Hegel.Connection (RequestError (..), sendRequestAndWait)

-- ----------------------------------------------------------------------------
-- Span label constants
-- ----------------------------------------------------------------------------

-- | Span label for list generation.
labelList :: Int
labelList = 1

-- | Span label for individual list elements.
labelListElement :: Int
labelListElement = 2

-- | Span label for set generation.
labelSet :: Int
labelSet = 3

-- | Span label for individual set elements.
labelSetElement :: Int
labelSetElement = 4

-- | Span label for map (dictionary) generation.
labelMap :: Int
labelMap = 5

-- | Span label for individual map entries.
labelMapEntry :: Int
labelMapEntry = 6

-- | Span label for tuple generation.
labelTuple :: Int
labelTuple = 7

-- | Span label for one-of dispatch.
labelOneOf :: Int
labelOneOf = 8

-- | Span label for optional generation.
labelOptional :: Int
labelOptional = 9

-- | Span label for fixed-dict generation.
labelFixedDict :: Int
labelFixedDict = 10

-- | Span label for flat-map generation.
labelFlatMap :: Int
labelFlatMap = 11

-- | Span label for filter generation.
labelFilter :: Int
labelFilter = 12

-- | Span label for mapped generation.
labelMapped :: Int
labelMapped = 13

-- | Span label for sampled-from generation.
labelSampledFrom :: Int
labelSampledFrom = 14

-- | Span label for enum variant generation.
labelEnumVariant :: Int
labelEnumVariant = 15

-- | Maximum number of filter attempts before calling @assume False@.
maxFilterAttempts :: Int
maxFilterAttempts = 3

-- ----------------------------------------------------------------------------
-- Generator GADT
-- ----------------------------------------------------------------------------

-- | A generator that produces values of type @a@.
--
-- The GADT constructors capture different generation strategies. 'Basic'
-- generators carry a CBOR schema sent to the server; the server returns a raw
-- 'Term' that is decoded by the transform. Other constructors compose
-- generators with client-side logic.
data Generator a where
  -- | A server-side generator with a CBOR schema and a client-side transform.
  Basic :: Term -> (Term -> a) -> Generator a
  -- | A mapped generator: apply @f@ to the result of the source.
  Mapped :: Generator b -> (b -> a) -> Generator a
  -- | A dependent generator: generate from source, then use the result to
  -- pick a second generator.
  FlatMapped :: Generator b -> (b -> Generator a) -> Generator a
  -- | A filtered generator: retry the source until the predicate holds.
  Filtered :: Generator a -> (a -> Bool) -> Generator a
  -- | A composite list generator using the collection protocol.
  CompositeList :: Generator a -> Int -> Maybe Int -> Generator [a]
  -- | A composite generator that runs an arbitrary IO action inside a
  -- labelled span.
  Composite :: Int -> (TestCase -> IO a) -> Generator a

instance Functor Generator where
  fmap f (Basic s t) = Basic s (f . t)  -- preserves schema
  fmap f other       = Mapped other f

-- ----------------------------------------------------------------------------
-- Span helpers
-- ----------------------------------------------------------------------------

-- | Run an IO action inside a span with the given label.
-- The span is stopped with @discard=False@ regardless of whether the action
-- throws.
group :: Int -> TestCase -> IO a -> IO a
group label tc action = do
  startSpan label tc
  result <- action `catch` \(e :: SomeException) -> do
    stopSpan False tc
    throwIO e
  stopSpan False tc
  return result

-- | Run an IO action inside a span with the given label. If the action
-- throws, the span is stopped with @discard=True@; otherwise @discard=False@.
discardableGroup :: Int -> TestCase -> IO a -> IO a
discardableGroup label tc action = do
  startSpan label tc
  result <- action `catch` \(e :: SomeException) -> do
    stopSpan True tc
    throwIO e
  stopSpan False tc
  return result

-- ----------------------------------------------------------------------------
-- Collection protocol
-- ----------------------------------------------------------------------------

-- | A server-side collection handle for generating variable-length sequences.
--
-- Collections communicate with the server to decide when to stop generating
-- elements. The 'collFinished' flag short-circuits subsequent 'collectionMore'
-- calls once the server signals completion.
data Collection = Collection
  { collFinished   :: !(IORef Bool)
  , collServerName :: !(IORef (Maybe Term))
  , collMinSize    :: !Int
  , collMaxSize    :: !(Maybe Int)
  }

-- | Create a new collection handle.
newCollection :: Int -> Maybe Int -> TestCase -> IO Collection
newCollection minSize maxSize _tc = do
  finished   <- newIORef False
  serverName <- newIORef Nothing
  return Collection
    { collFinished   = finished
    , collServerName = serverName
    , collMinSize    = minSize
    , collMaxSize    = maxSize
    }

-- | Send a command on the test case's data channel. Handles StopTest.
sendCommand :: TestCase -> Term -> IO Term
sendCommand tc msg = do
  aborted <- readIORef (tcTestAborted tc)
  if aborted
    then throwIO DataExhausted
    else sendRequestAndWait (tcChannel tc) msg
           `catch` \(e :: RequestError) ->
             if reqErrorType e == "StopTest"
               then do
                 writeIORef (tcTestAborted tc) True
                 throwIO DataExhausted
               else throwIO e

-- | Lazily initialise the server-side collection and return its handle.
getServerName :: Collection -> TestCase -> IO Term
getServerName coll tc = do
  mName <- readIORef (collServerName coll)
  case mName of
    Just name -> return name
    Nothing -> do
      let maxSizeVal = maybe TNull TInt (collMaxSize coll)
      let cmd = TMap
            [ (TString "command",  TString "new_collection")
            , (TString "min_size", TInt (collMinSize coll))
            , (TString "max_size", maxSizeVal)
            ]
      result <- sendCommand tc cmd
      writeIORef (collServerName coll) (Just result)
      return result

-- | Query the server for whether more elements should be generated.
-- Once this returns 'False', subsequent calls return 'False' immediately.
collectionMore :: Collection -> TestCase -> IO Bool
collectionMore coll tc = do
  done <- readIORef (collFinished coll)
  if done
    then return False
    else do
      serverName <- getServerName coll tc
      let cmd = TMap
            [ (TString "command",    TString "collection_more")
            , (TString "collection", serverName)
            ]
      result <- sendCommand tc cmd
      let more = case result of
            TBool b -> b
            _       -> error "collectionMore: expected boolean from server"
      unless more $ writeIORef (collFinished coll) True
      return more

-- | Reject the last element of the collection. No-op if already finished.
collectionReject :: Collection -> TestCase -> IO ()
collectionReject coll tc = do
  done <- readIORef (collFinished coll)
  if done
    then return ()
    else do
      serverName <- getServerName coll tc
      let cmd = TMap
            [ (TString "command",    TString "collection_reject")
            , (TString "collection", serverName)
            ]
      _ <- sendCommand tc cmd
      return ()

-- ----------------------------------------------------------------------------
-- Draw (evaluation)
-- ----------------------------------------------------------------------------

-- | Produce a typed value from a generator using the given test case.
draw :: TestCase -> Generator a -> IO a
draw tc gen = case gen of
  Basic schema transform -> do
    raw <- generateFromSchema schema tc
    return (transform raw)

  Mapped source f ->
    group labelMapped tc $ do
      v <- draw tc source
      return (f v)

  FlatMapped source f ->
    discardableGroup labelFlatMap tc $ do
      v <- draw tc source
      draw tc (f v)

  Filtered source predicate ->
    let attempt i
          | i > maxFilterAttempts = do { assume tc False; error "unreachable" }
          | otherwise = do
              startSpan labelFilter tc
              v <- draw tc source
              if predicate v
                then do
                  stopSpan False tc
                  return v
                else do
                  stopSpan True tc
                  attempt (i + 1)
    in attempt 1

  CompositeList elements minSize maxSize ->
    group labelList tc $ do
      coll <- newCollection minSize maxSize tc
      let collect acc = do
            more <- collectionMore coll tc
            if more
              then do
                v <- draw tc elements
                collect (v : acc)
              else return (reverse acc)
      collect []

  Composite label generateFn ->
    group label tc (generateFn tc)

-- ----------------------------------------------------------------------------
-- Combinator functions
-- ----------------------------------------------------------------------------

-- | Map a function over a generator's output.
--
-- When the generator is 'Basic', the schema is preserved and transforms are
-- composed. Otherwise a 'Mapped' wrapper is created.
gmap :: (a -> b) -> Generator a -> Generator b
gmap f (Basic schema transform) = Basic schema (f . transform)
gmap f other                    = Mapped other f

-- | Create a dependent generator. The function receives the generated value
-- and returns a new generator whose output is the final result.
flatMap :: (a -> Generator b) -> Generator a -> Generator b
flatMap f gen = FlatMapped gen f

-- | Filter a generator's output with a predicate. Retries up to
-- 'maxFilterAttempts' times; if all attempts fail, @assume False@ is called.
gfilter :: (a -> Bool) -> Generator a -> Generator a
gfilter predicate gen = Filtered gen predicate

-- ----------------------------------------------------------------------------
-- Introspection
-- ----------------------------------------------------------------------------

-- | Returns 'True' if the generator is a 'Basic' generator.
isBasic :: Generator a -> Bool
isBasic (Basic _ _) = True
isBasic _           = False

-- | Returns the schema and transform of a 'Basic' generator, or 'Nothing'.
asBasic :: Generator a -> Maybe (Term, Term -> a)
asBasic (Basic schema transform) = Just (schema, transform)
asBasic _                        = Nothing