Unwrap.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

{-|
Module      :  Data.Aeson.Schema.TH.Unwrap
Maintainer  :  Brandon Chinn <brandonchinn178@gmail.com>
Stability   :  experimental
Portability :  portable

The 'unwrap' quasiquoter.
-}
module Data.Aeson.Schema.TH.Unwrap where

import Control.Monad ((<=<), (>=>))
import Data.Bifunctor (first)
import qualified Data.List.NonEmpty as NonEmpty
import Language.Haskell.TH
import Language.Haskell.TH.Quote (QuasiQuoter (..))

import Data.Aeson.Schema.Internal (Object, SchemaResult)
import Data.Aeson.Schema.Key (fromSchemaKeyV)
import Data.Aeson.Schema.TH.Parse (
  GetterOperation (..),
  GetterOps,
  UnwrapSchema (..),
  parseUnwrapSchema,
 )
import Data.Aeson.Schema.TH.Utils (
  reifySchema,
  resolveSchemaType,
  schemaTypeVToTypeQ,
  schemaVToTypeQ,
 )
import Data.Aeson.Schema.Type (
  Schema' (..),
  SchemaType' (..),
  SchemaTypeV,
  SchemaV,
  showSchemaTypeV,
  toSchemaObjectV,
 )

-- | Defines a QuasiQuoter to extract a schema within the given schema.
--
--  The base schema needs to be defined in a separate module.
--
--  For example:
--
--  > -- | MyFoo ~ Object [schema| { b: Maybe Bool } |]
--  > type MyFoo = [unwrap| MySchema.foo.nodes[] |]
--
--  If the schema is imported qualified, you can use parentheses to distinguish it from the
--  expression:
--
--  > type MyFoo = [unwrap| (MyModule.Schema).foo.nodes[] |]
--
--  You can then use the type alias as usual:
--
--  > parseBar :: MyFoo -> String
--  > parseBar = maybe "null" show . [get| .b |]
--  >
--  > foo = map parseBar [get| result.foo.nodes[] |]
--
--  The syntax is mostly the same as 'Data.Aeson.Schema.TH.get', except the operations run on the
--  type itself, instead of the values. Differences from 'Data.Aeson.Schema.TH.get':
--
--  * @x!@ is only valid if @x@ is a @Maybe a@ type. Returns @a@, the type wrapped in the 'Maybe'.
--
--  * @x?@ is the same as @x!@.
--
--  * @x[]@ is only valid if @x@ is a @[a]@ type. Returns @a@, the type contained in the list.
--
--  * @x\@#@ is only valid if @x@ is a @SumType@. Returns the type at that branch in the sum type.
unwrap :: QuasiQuoter
unwrap =
  QuasiQuoter
    { quoteExp = error "Cannot use `unwrap` for Exp"
    , quoteDec = error "Cannot use `unwrap` for Dec"
    , quoteType = parseUnwrapSchema >=> generateUnwrapSchema
    , quotePat = error "Cannot use `unwrap` for Pat"
    }

generateUnwrapSchema :: UnwrapSchema -> TypeQ
generateUnwrapSchema UnwrapSchema{..} = reifySchema startSchema >>= unwrapSchema getterOps

-- | Unwrap the given schema by applying the given operations, stripping out functors.
unwrapSchema :: GetterOps -> SchemaV -> TypeQ
unwrapSchema = unwrapSchemaUsing StripFunctors

-- | Unwrap the given schema by applying the given operations, using the given 'FunctorHandler'.
unwrapSchemaUsing :: FunctorHandler -> GetterOps -> SchemaV -> TypeQ
unwrapSchemaUsing functorHandler getterOps = toResultTypeQ <=< flip go (NonEmpty.toList getterOps) . toSchemaObjectV
  where
    toResultTypeQ :: UnwrapSchemaResult -> TypeQ
    toResultTypeQ = \case
      -- special case SchemaObject to make it further inspectable
      SchemaResult (SchemaObject pairs) -> [t|Object $(schemaVToTypeQ (Schema pairs))|]
      SchemaResult schemaType -> [t|SchemaResult $(schemaTypeVToTypeQ schemaType)|]
      SchemaResultList schemaResult -> appT listT (toResultTypeQ schemaResult)
      SchemaResultTuple schemaResults -> foldl appT (tupleT $ length schemaResults) $ map toResultTypeQ schemaResults
      SchemaResultWrapped functorTy schemaResult ->
        let handleFunctor ty =
              case functorHandler of
                ApplyFunctors -> AppT functorTy ty
                StripFunctors -> ty
         in handleFunctor <$> toResultTypeQ schemaResult

    go :: SchemaTypeV -> [GetterOperation] -> Q UnwrapSchemaResult
    go schemaType [] = pure $ SchemaResult schemaType
    go schemaType' (op : ops) = do
      schemaType <- resolveSchemaType schemaType'

      let invalid message = fail $ message ++ ": " ++ showSchemaTypeV schemaType
          wrapMaybe = SchemaResultWrapped (ConT ''Maybe)
          wrapList = SchemaResultWrapped ListT

      case op of
        GetterKey key ->
          case schemaType of
            SchemaObject pairs ->
              case lookup key $ map (first fromSchemaKeyV) pairs of
                Just inner -> go inner ops
                Nothing -> invalid $ "Key '" ++ key ++ "' does not exist in schema"
            _ -> invalid $ "Cannot get key '" ++ key ++ "' in schema"
        GetterBang ->
          case schemaType of
            SchemaMaybe inner -> go inner ops
            SchemaTry inner -> go inner ops
            _ -> invalid "Cannot use `!` operator on schema"
        GetterMapMaybe ->
          case schemaType of
            SchemaMaybe inner -> wrapMaybe <$> go inner ops
            SchemaTry inner -> wrapMaybe <$> go inner ops
            _ -> invalid "Cannot use `?` operator on schema"
        GetterMapList ->
          case schemaType of
            SchemaList inner -> wrapList <$> go inner ops
            _ -> invalid "Cannot use `[]` operator on schema"
        GetterBranch branch ->
          case schemaType of
            SchemaUnion schemas ->
              if branch < length schemas
                then go (schemas !! branch) ops
                else invalid "Branch out of bounds for schema"
            _ -> invalid "Cannot use `@` operator on schema"
        -- suffixes; ops should be empty

        GetterList elemOps ->
          case schemaType of
            SchemaObject _ -> do
              elemSchemas <- traverse (go schemaType . NonEmpty.toList) elemOps
              let elemSchema = NonEmpty.head elemSchemas
              if all (== elemSchema) elemSchemas
                then pure $ SchemaResultList elemSchema
                else invalid "List contains different types in schema"
            _ -> invalid "Cannot get keys in schema"
        GetterTuple elemOps ->
          case schemaType of
            SchemaObject _ -> SchemaResultTuple <$> mapM (go schemaType . NonEmpty.toList) (NonEmpty.toList elemOps)
            _ -> invalid "Cannot get keys in schema"

data UnwrapSchemaResult
  = SchemaResult SchemaTypeV
  | SchemaResultList UnwrapSchemaResult
  | SchemaResultTuple [UnwrapSchemaResult]
  | -- | Type should be of kind `* -> *`
    SchemaResultWrapped Type UnwrapSchemaResult
  deriving (Eq)

-- | A data type that indicates how to handle functors when unwrapping a schema.
data FunctorHandler
  = -- | handleFunctor Maybe Int ==> Maybe Int
    ApplyFunctors
  | -- | handleFunctor Maybe Int ==> Int
    StripFunctors