Collections.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Collection generators: lists and hashmaps.
--
-- These generators handle variable-length collections. When all element
-- generators are 'Basic', a single schema is sent to the server for efficient
-- generation. When elements are non-basic (e.g. filtered or flat-mapped), the
-- collection protocol is used to generate elements one at a time.
module Hegel.Generator.Collections
  ( lists
  , hashmaps
  ) where

import Codec.CBOR.Term (Term (..))
import Data.Maybe (catMaybes)

import Hegel.Generator (Generator (..), asBasic)
import Hegel.Generator.Range (RangeOpts, validateSizeBounds)

-- ----------------------------------------------------------------------------
-- Lists
-- ----------------------------------------------------------------------------

-- | Generator for lists of elements.
--
-- Two code paths:
--
-- 1. If the element generator is 'Basic': sends a @list@ schema to the server
--    and lets it generate the entire list in one round-trip. The element
--    transform is lifted to apply to every item.
--
-- 2. If the element generator is non-basic: uses 'CompositeList' with the
--    collection protocol to generate elements one at a time.
lists :: Generator a -> RangeOpts Int -> Generator [a]
lists elements opts =
  let (ms, maxSize) = validateSizeBounds "lists" opts
  in
  case asBasic elements of
    Just (elemSchema, elemTransform) ->
      let pairs = catMaybes
            [ Just (TString "type",     TString "list")
            , Just (TString "elements", elemSchema)
            , Just (TString "min_size", TInt ms)
            , fmap (\mx -> (TString "max_size", TInt mx)) maxSize
            ]
          listTransform raw = case raw of
            TList items -> map elemTransform items
            _           -> error "lists: expected array from server"
      in Basic (TMap pairs) listTransform

    Nothing ->
      CompositeList elements ms maxSize

-- ----------------------------------------------------------------------------
-- Hashmaps
-- ----------------------------------------------------------------------------

-- | Generator for dictionaries (hash maps) as lists of key-value pairs.
--
-- Both the key and value generators must be 'Basic'. The server returns the
-- dict as a list of @[key, value]@ pairs, which are transformed to Haskell
-- tuples.
--
-- Schema: @{\"type\": \"dict\", \"keys\": kSchema, \"values\": vSchema,
-- \"min_size\": N, \"max_size\"?: N}@
hashmaps :: Generator k -> Generator v -> RangeOpts Int -> Generator [(k, v)]
hashmaps keys values opts =
  let (ms, maxSize) = validateSizeBounds "hashmaps" opts in
  let (keySchema, keyTransform) = case asBasic keys of
        Just (s, t) -> (s, t)
        Nothing     -> error "hashmaps: keys generator must be a Basic generator"
      (valSchema, valTransform) = case asBasic values of
        Just (s, t) -> (s, t)
        Nothing     -> error "hashmaps: values generator must be a Basic generator"
      pairs = catMaybes
        [ Just (TString "type",     TString "dict")
        , Just (TString "keys",     keySchema)
        , Just (TString "values",   valSchema)
        , Just (TString "min_size", TInt ms)
        , fmap (\mx -> (TString "max_size", TInt mx)) maxSize
        ]
      transform raw = case raw of
        TList kvPairs -> map transformPair kvPairs
        _             -> error "hashmaps: expected array from server"
      transformPair pair = case pair of
        TList [k, v] -> (keyTransform k, valTransform v)
        _            -> error "hashmaps: expected [k, v] pair from server"
  in Basic (TMap pairs) transform