Schema.hs

{-|
Module      :  Data.Aeson.Schema.TH.Schema
Maintainer  :  Brandon Chinn <brandon@leapyear.io>
Stability   :  experimental
Portability :  portable

The 'schema' quasiquoter.
-}
{-# 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
    ( Schema'(..)
    , SchemaObjectMapV
    , SchemaType'(..)
    , SchemaTypeV
    , fromSchemaV
    , showSchemaTypeV
    )
import Data.Aeson.Schema.Utils.Invariant (unreachable)
import Data.Aeson.Schema.Utils.NameLike (NameLike(..))

-- | Defines a QuasiQuoter for writing schemas.
--
-- Example:
--
-- > import Data.Aeson.Schema (schema)
-- >
-- > type MySchema = [schema|
-- >   {
-- >     foo: {
-- >       a: Int,
-- >       // you can add comments like this
-- >       nodes: List {
-- >         b: Maybe Bool,
-- >       },
-- >       c: Text,
-- >       d: Text,
-- >       e: MyType,
-- >       f: Maybe List {
-- >         name: Text,
-- >       },
-- >     },
-- >   }
-- > |]
--
-- Syntax:
--
-- * @{ key: \<schema\>, ... }@ corresponds to a JSON 'Data.Aeson.Schema.Object' with the given key
--   mapping to the given schema.
--
-- * @Bool@, @Int@, @Double@, and @Text@ correspond to the usual Haskell values.
--
-- * @Maybe \<schema\>@ and @List \<schema\>@ correspond to @Maybe@ and @[]@, containing values
--   specified by the provided schema (no parentheses needed).
--
-- * @Try \<schema\>@ corresponds to @Maybe@, where the value will be @Just@ if the given schema
--   successfully parses the value, or @Nothing@ otherwise. Different from @Maybe \<schema\>@,
--   where parsing @{ "foo": true }@ with @{ foo: Try Int }@ returns @Nothing@, whereas it would
--   be a parse error with @{ foo: Maybe Int }@ (added in v1.2.0)
--
-- * Any other uppercase identifier corresponds to the respective type in scope -- requires a
--   FromJSON instance.
--
-- Advanced syntax:
--
-- * @\<schema1\> | \<schema2\>@ corresponds to a JSON value that matches one of the given schemas.
--   When extracted from an 'Data.Aeson.Schema.Object', it deserializes into a
--   'Data.Aeson.Schema.Utils.Sum.JSONSum' object. (added in v1.1.0)
--
-- * @{ [key]: \<schema\> }@ uses the current object to resolve the keys in the given schema. Only
--   object schemas are allowed here. (added in v1.2.0)
--
-- * @{ key: #Other, ... }@ maps the given key to the @Other@ schema. The @Other@ schema needs to
--   be defined in another module.
--
-- * @{ #Other, ... }@ extends this schema with the @Other@ schema. The @Other@ schema needs to
--   be defined in another module.
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

-- | Get the SchemaObjectMapV for the given SchemaDefObjItem, along with where the SchemaObjectMapV
-- came from.
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
    -- should return true if it's at all possible to get a valid parse
    isValidPhantomSchema = \case
      SchemaMaybe inner -> isValidPhantomSchema inner
      SchemaTry _ -> True -- even if inner is a non-object schema, it'll still parse to be Nothing
      SchemaUnion schemas -> any isValidPhantomSchema schemas
      SchemaObject _ -> True
      SchemaInclude _ -> True
      _ -> False

-- | Resolve the given keys with the following rules:
--
-- 1. Any explicitly provided keys shadow/overwrite imported keys
-- 2. Fail if duplicate keys are both explicitly provided
-- 3. Fail if duplicate keys are both imported
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)

{- SchemaDef conversions -}

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 -> return $ SchemaInclude $ Left $ NameRef other
  SchemaDefUnion schemas -> SchemaUnion . NonEmpty.toList <$> mapM fromSchemaDefType schemas
  SchemaDefObj items     -> SchemaObject <$> generateSchemaObjectV items

{- LookupMap utilities -}

type LookupMap k v = [(k, v)]

-- | Distribute the given element across the values in the map.
distribute :: LookupMap k v -> a -> LookupMap k (a, v)
distribute lookupMap a = map (fmap (a,)) lookupMap

-- | Find all values with the same key (according to the given function) and group them.
--
-- Invariants:
-- * [v] has length > 0
-- * If the first occurence of k1 is before the first occurence of k2, k1 is before k2
--   in the result
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
    -- don't use sort; keys should stay in the same order
    distinctKeys = nubBy ((==) `on` f) $ map fst pairs

    groups = HashMap.fromListWith (flip (++)) $ map (\(k, v) -> (f k, [v])) pairs