Combinators.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Combinator generators: sampling, alternation, optionality, IP addresses,
-- and tuples.
--
-- These generators compose primitive and collection generators into richer
-- generation strategies. When all constituent generators are 'Basic', a single
-- compound schema is sent to the server. When any generator is non-basic, a
-- 'Composite' wrapper is used with client-side orchestration.
module Hegel.Generator.Combinators
  ( sampledFrom
  , oneOf
  , optional
  , ipAddresses
  , tuples2
  , tuples3
  , tuples4
  ) where

import Codec.CBOR.Term (Term (..))
import Data.Array (listArray, (!))
import Data.Default (def)
import Data.Maybe (catMaybes, isJust)
import Data.Text (Text)

import Hegel.Client (TestCase, generateFromSchema)
import Hegel.Generator
  ( Generator (..)
  , asBasic
  , draw
  , labelOneOf
  , labelTuple
  )
import Hegel.Generator.Primitives
  ( RangeOpts (..)
  , integers
  , ipv4Addresses
  , ipv6Addresses
  , just
  )

-- ----------------------------------------------------------------------------
-- CBOR extraction helper
-- ----------------------------------------------------------------------------

-- | Extract an 'Int' from a CBOR term.
extractInt :: Term -> Int
extractInt (TInt n)     = n
extractInt (TInteger n) = fromIntegral n
extractInt t            = error $ "extractInt: expected TInt/TInteger, got " ++ show t

-- ----------------------------------------------------------------------------
-- sampledFrom
-- ----------------------------------------------------------------------------

-- | Generator that samples uniformly from a non-empty list of values.
--
-- Implemented as an integer index generator: picks an index in @[0, n-1]@ and
-- returns the corresponding element from an array.
sampledFrom :: [a] -> Generator a
sampledFrom [] = error "sampledFrom: requires at least one element"
sampledFrom options =
  let n   = length options
      arr = listArray (0, n - 1) options
  in fmap (arr !) (integers def
      { rangeMin = Just 0
      , rangeMax = Just (n - 1)
      })

-- ----------------------------------------------------------------------------
-- oneOf
-- ----------------------------------------------------------------------------

-- | Generator that picks from one of the given generators.
--
-- Two code paths:
--
-- 1. All generators are 'Basic': uses a tagged-tuple dispatch schema.
--    Each alternative is wrapped as @{\"type\": \"tuple\", \"elements\":
--    [{\"const\": tag}, schema]}@. The server returns @[tag, value]@ and the
--    client dispatches to the correct transform.
--
-- 2. Any generator is non-basic: generates an index inside a 'labelOneOf'
--    span, then delegates to the selected generator.
oneOf :: [Generator a] -> Generator a
oneOf [] = error "oneOf: requires at least one generator"
oneOf generators =
  let basics   = map asBasic generators
      allBasic = all isJust basics
  in
  if not allBasic
    then
      -- Composite path: generate index, then delegate
      let gens = listArray (0, n - 1) generators
      in Composite labelOneOf $ \tc -> do
           idx <- drawIndex n tc
           draw tc (gens ! idx)
    else
      -- Tagged-tuple path
      let justBasics = catMaybes basics
          taggedSchemas = zipWith mkTaggedSchema [0..] (map fst justBasics)
          transforms    = listArray (0, length justBasics - 1) (map snd justBasics)
          dispatch raw = case raw of
            TList [tag, value] ->
              let idx = extractInt tag
              in (transforms ! idx) value
            _ -> error "oneOf: expected [tag, value] tuple from server"
      in Basic
           (TMap [(TString "one_of", TList taggedSchemas)])
           dispatch
  where
    n = length generators

    drawIndex :: Int -> TestCase -> IO Int
    drawIndex count tc = do
      let schema = TMap
            [ (TString "type",      TString "integer")
            , (TString "min_value", TInt 0)
            , (TString "max_value", TInt (count - 1))
            ]
      raw <- generateFromSchema schema tc
      return (extractInt raw)

    mkTaggedSchema :: Int -> Term -> Term
    mkTaggedSchema i elemSchema = TMap
      [ (TString "type",     TString "tuple")
      , (TString "elements", TList
          [ TMap [(TString "const", TInt i)]
          , elemSchema
          ])
      ]

-- ----------------------------------------------------------------------------
-- optional
-- ----------------------------------------------------------------------------

-- | Generator that produces either 'Nothing' or @'Just' value@.
--
-- Equivalent to @oneOf [just Nothing, fmap Just gen]@.
optional :: Generator a -> Generator (Maybe a)
optional gen = oneOf [just Nothing, fmap Just gen]

-- ----------------------------------------------------------------------------
-- IP addresses
-- ----------------------------------------------------------------------------

-- | Generator for IP address strings.
--
-- * @Nothing@: generates either IPv4 or IPv6 addresses.
-- * @Just 4@: generates IPv4 addresses only.
-- * @Just 6@: generates IPv6 addresses only.
ipAddresses :: Maybe Int -> Generator Text
ipAddresses Nothing  = oneOf [ipv4Addresses, ipv6Addresses]
ipAddresses (Just 4) = ipv4Addresses
ipAddresses (Just 6) = ipv6Addresses
ipAddresses (Just v) = error $ "ipAddresses: invalid version " ++ show v

-- ----------------------------------------------------------------------------
-- Tuples
-- ----------------------------------------------------------------------------

-- | Generator for 2-element tuples.
--
-- When both elements are 'Basic', a single @tuple@ schema is sent.
-- Otherwise, elements are generated separately inside a 'labelTuple' span.
tuples2 :: Generator a -> Generator b -> Generator (a, b)
tuples2 g1 g2 =
  case (asBasic g1, asBasic g2) of
    (Just (s1, t1), Just (s2, t2)) ->
      let schema = TMap
            [ (TString "type",     TString "tuple")
            , (TString "elements", TList [s1, s2])
            ]
          transform raw = case raw of
            TList [v1, v2] -> (t1 v1, t2 v2)
            _ -> error "tuples2: expected 2-element array from server"
      in Basic schema transform
    _ ->
      Composite labelTuple $ \tc -> do
        a <- draw tc g1
        b <- draw tc g2
        return (a, b)

-- | Generator for 3-element tuples.
--
-- When all elements are 'Basic', a single @tuple@ schema is sent.
-- Otherwise, elements are generated separately inside a 'labelTuple' span.
tuples3 :: Generator a -> Generator b -> Generator c -> Generator (a, b, c)
tuples3 g1 g2 g3 =
  case (asBasic g1, asBasic g2, asBasic g3) of
    (Just (s1, t1), Just (s2, t2), Just (s3, t3)) ->
      let schema = TMap
            [ (TString "type",     TString "tuple")
            , (TString "elements", TList [s1, s2, s3])
            ]
          transform raw = case raw of
            TList [v1, v2, v3] -> (t1 v1, t2 v2, t3 v3)
            _ -> error "tuples3: expected 3-element array from server"
      in Basic schema transform
    _ ->
      Composite labelTuple $ \tc -> do
        a <- draw tc g1
        b <- draw tc g2
        c <- draw tc g3
        return (a, b, c)

-- | Generator for 4-element tuples.
--
-- When all elements are 'Basic', a single @tuple@ schema is sent.
-- Otherwise, elements are generated separately inside a 'labelTuple' span.
tuples4 :: Generator a -> Generator b -> Generator c -> Generator d
        -> Generator (a, b, c, d)
tuples4 g1 g2 g3 g4 =
  case (asBasic g1, asBasic g2, asBasic g3, asBasic g4) of
    (Just (s1, t1), Just (s2, t2), Just (s3, t3), Just (s4, t4)) ->
      let schema = TMap
            [ (TString "type",     TString "tuple")
            , (TString "elements", TList [s1, s2, s3, s4])
            ]
          transform raw = case raw of
            TList [v1, v2, v3, v4] -> (t1 v1, t2 v2, t3 v3, t4 v4)
            _ -> error "tuples4: expected 4-element array from server"
      in Basic schema transform
    _ ->
      Composite labelTuple $ \tc -> do
        a <- draw tc g1
        b <- draw tc g2
        c <- draw tc g3
        d <- draw tc g4
        return (a, b, c, d)