{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Aeson.Schema.TH.Schema (schema) where
import Control.Monad (unless, (>=>))
import Data.Function (on)
import Data.Hashable (Hashable)
import qualified Data.HashMap.Strict as HashMap
import Data.List (nubBy)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import Language.Haskell.TH
import Language.Haskell.TH.Quote (QuasiQuoter(..))
import Data.Aeson.Schema.Key (SchemaKey'(..), SchemaKeyV, fromSchemaKeyV)
import Data.Aeson.Schema.TH.Parse
(SchemaDef(..), SchemaDefObjItem(..), SchemaDefObjKey(..), parseSchemaDef)
import Data.Aeson.Schema.TH.Utils (reifySchema, schemaVToTypeQ)
import Data.Aeson.Schema.Type
( NameLike(..)
, Schema'(..)
, SchemaObjectMapV
, SchemaType'(..)
, SchemaTypeV
, fromSchemaV
, showSchemaTypeV
, toSchemaObjectV
)
import Data.Aeson.Schema.Utils.Invariant (unreachable)
schema :: QuasiQuoter
schema = QuasiQuoter
{ quoteExp = error "Cannot use `schema` for Exp"
, quoteDec = error "Cannot use `schema` for Dec"
, quoteType = parseSchemaDef >=> \case
SchemaDefObj items -> generateSchemaObject items
_ -> fail "`schema` definition must be an object"
, quotePat = error "Cannot use `schema` for Pat"
}
where
generateSchemaObject items = schemaVToTypeQ . Schema =<< generateSchemaObjectV items
data KeySource = Provided | Imported
deriving (Show, Eq)
generateSchemaObjectV :: NonEmpty SchemaDefObjItem -> Q SchemaObjectMapV
generateSchemaObjectV schemaDefObjItems = do
schemaObjectMapsWithSource <- mapM getSchemaObjectMap schemaDefObjItems
let schemaObjectMaps :: LookupMap SchemaKeyV (KeySource, SchemaTypeV)
schemaObjectMaps = concatMap (uncurry distribute) schemaObjectMapsWithSource
either fail return $ resolveKeys schemaObjectMaps
getSchemaObjectMap :: SchemaDefObjItem -> Q (SchemaObjectMapV, KeySource)
getSchemaObjectMap = \case
SchemaDefObjPair (schemaDefKey, schemaDefType) -> do
let schemaKey = fromSchemaDefKey schemaDefKey
schemaType <- fromSchemaDefType schemaDefType
case schemaKey of
PhantomKey _ ->
unless (isValidPhantomSchema schemaType) $
fail $ "Invalid schema for '" ++ fromSchemaKeyV schemaKey ++ "': " ++ showSchemaTypeV schemaType
_ -> return ()
return ([(schemaKey, schemaType)], Provided)
SchemaDefObjExtend other -> do
schemaV <- reifySchema other
return (fromSchemaV schemaV, Imported)
where
isValidPhantomSchema = \case
SchemaMaybe inner -> isValidPhantomSchema inner
SchemaTry _ -> True
SchemaUnion schemas -> any isValidPhantomSchema schemas
SchemaObject _ -> True
_ -> False
resolveKeys :: Show a => LookupMap SchemaKeyV (KeySource, a) -> Either String (LookupMap SchemaKeyV a)
resolveKeys = mapM (uncurry resolveKey) . groupByKeyWith fromSchemaKeyV
where
resolveKey key sourcesAndVals =
let filterSource source = map snd $ filter ((== source) . fst) sourcesAndVals
numProvided = length $ filterSource Provided
numImported = length $ filterSource Imported
in if
| numProvided > 1 -> Left $ "Key '" ++ fromSchemaKeyV key ++ "' specified multiple times"
| [val] <- filterSource Provided -> Right (key, val)
| numImported > 1 -> Left $ "Key '" ++ fromSchemaKeyV key ++ "' declared in multiple imported schemas"
| [val] <- filterSource Imported -> Right (key, val)
| otherwise -> unreachable $ "resolveKey received: " ++ show (key, sourcesAndVals)
fromSchemaDefKey :: SchemaDefObjKey -> SchemaKeyV
fromSchemaDefKey = \case
SchemaDefObjKeyNormal key -> NormalKey key
SchemaDefObjKeyPhantom key -> PhantomKey key
fromSchemaDefType :: SchemaDef -> Q SchemaTypeV
fromSchemaDefType = \case
SchemaDefType name -> return $ SchemaScalar $ NameRef name
SchemaDefMaybe inner -> SchemaMaybe <$> fromSchemaDefType inner
SchemaDefTry inner -> SchemaTry <$> fromSchemaDefType inner
SchemaDefList inner -> SchemaList <$> fromSchemaDefType inner
SchemaDefInclude other -> toSchemaObjectV <$> reifySchema other
SchemaDefUnion schemas -> SchemaUnion . NonEmpty.toList <$> mapM fromSchemaDefType schemas
SchemaDefObj items -> SchemaObject <$> generateSchemaObjectV items
type LookupMap k v = [(k, v)]
distribute :: LookupMap k v -> a -> LookupMap k (a, v)
distribute lookupMap a = map (fmap (a,)) lookupMap
groupByKeyWith :: (Eq a, Hashable a) => (k -> a) -> LookupMap k v -> LookupMap k [v]
groupByKeyWith f pairs = map (\key -> (key, groups HashMap.! f key)) distinctKeys
where
distinctKeys = nubBy ((==) `on` f) $ map fst pairs
groups = HashMap.fromListWith (flip (++)) $ map (\(k, v) -> (f k, [v])) pairs