Schemas.hs

{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

{-| This module contains the implementation of the @dhall rewrite-with-schemas@
    subcommand
-}

module Dhall.Schemas
    ( -- | Schemas
      schemasCommand
    , Schemas(..)
    , rewriteWithSchemas
    , SchemasError(..)
    ) where

import Control.Applicative (empty)
import Control.Exception   (Exception)
import Data.Maybe          (fromMaybe)
import Data.Text           (Text)
import Data.Void           (Void)
import Dhall.Crypto        (SHA256Digest)
import Dhall.Map           (Map)
import Dhall.Pretty        (CharacterSet (..), detectCharacterSet)
import Dhall.Src           (Src)
import Dhall.Syntax        (Expr (..), Import, Var (..))
import Dhall.Util
    ( Censor (..)
    , Header (..)
    , Input (..)
    , MultipleCheckFailed (..)
    , OutputMode (..)
    )

import qualified Control.Exception                  as Exception
import qualified Data.Map
import qualified Data.Maybe                         as Maybe
import qualified Data.Text.IO                       as Text.IO
import qualified Data.Void                          as Void
import qualified Dhall.Core                         as Core
import qualified Dhall.Import                       as Import
import qualified Dhall.Map                          as Map
import qualified Dhall.Normalize                    as Normalize
import qualified Dhall.Parser                       as Parser
import qualified Dhall.Pretty
import qualified Dhall.Substitution                 as Substitution
import qualified Dhall.Syntax                       as Syntax
import qualified Dhall.TypeCheck                    as TypeCheck
import qualified Dhall.Util                         as Util
import qualified Lens.Micro                         as Lens
import qualified Prettyprinter                      as Pretty
import qualified Prettyprinter.Render.Terminal      as Pretty.Terminal
import qualified Prettyprinter.Render.Text          as Pretty.Text
import qualified System.AtomicWrite.Writer.LazyText as AtomicWrite
import qualified System.Console.ANSI                as ANSI
import qualified System.IO                          as IO

-- | Arguments to the @rewrite-with-schemas@ subcommand
data Schemas = Schemas
    { chosenCharacterSet :: Maybe CharacterSet
    , censor             :: Censor
    , input              :: Input
    , outputMode         :: OutputMode
    , schemas            :: Text
    }

-- | Implementation of the @dhall rewrite-with-schemas@ subcommand
schemasCommand :: Schemas -> IO ()
schemasCommand Schemas{..} = do
    (inputName, originalText) <- case input of
        InputFile file -> (,) file <$> Text.IO.readFile file
        StandardInput  -> (,) "(input)" <$> Text.IO.getContents

    (Header header, expression) <- Util.getExpressionAndHeaderFromStdinText censor inputName originalText

    let characterSet = fromMaybe (detectCharacterSet expression) chosenCharacterSet

    schemasRecord <- Core.throws (Parser.exprFromText "(schemas)" schemas)

    schemasExpression <- rewriteWithSchemas schemasRecord expression

    let docStream =
            Dhall.Pretty.layout
                (   Pretty.pretty header
                <>  Dhall.Pretty.prettyCharacterSet characterSet schemasExpression
                <>  "\n"
                )

    let schemasText = Pretty.Text.renderStrict docStream

    case outputMode of
        Write ->
            case input of
                InputFile file ->
                    if originalText == schemasText
                        then return ()
                        else AtomicWrite.atomicWriteFile
                                file
                                (Pretty.Text.renderLazy docStream)
                StandardInput -> do
                    supportsANSI <- ANSI.hSupportsANSI IO.stdout

                    Pretty.Terminal.renderIO
                        IO.stdout
                        (if supportsANSI
                            then fmap Dhall.Pretty.annToAnsiStyle docStream
                            else Pretty.unAnnotateS docStream)

        Check ->
            if originalText == schemasText
                then return ()
                else do
                    let command = "rewrite-with-schemas"

                    let modified = "rewritten"

                    let inputs = pure input

                    Exception.throwIO MultipleCheckFailed{..}

decodeSchema :: Expr s Void -> Maybe (Expr s Void, Map Text (Expr s Void))
decodeSchema (RecordLit m)
        | Just  _Type               <- Core.recordFieldValue <$> Map.lookup "Type" m
        , Just (RecordLit _default) <- Core.recordFieldValue <$> Map.lookup "default" m =
            Just (_Type, Core.recordFieldValue <$> _default)
decodeSchema _ =
    Nothing

decodeSchemas
    :: Expr s Void
    -> Maybe (Data.Map.Map SHA256Digest (Text, Map Text (Expr s Void)))
decodeSchemas (RecordLit keyValues) = do
    m <- traverse (decodeSchema . Core.recordFieldValue) keyValues

    let typeMetadata = Data.Map.fromList $ do
            (name, (_Type, _default)) <- Map.toList m

            return (Import.hashExpression (Syntax.denote _Type), (name, _default))

    return typeMetadata
decodeSchemas  _ =
    empty

-- | Simplify a Dhall expression using a record of schemas
rewriteWithSchemas
    :: Expr Src Import
    -- ^ Record of schemas
    -> Expr Src Import
    -- ^ Expression to simplify using the supplied schemas
    -> IO (Expr Src Import)
rewriteWithSchemas _schemas expression = do
    resolvedSchemas    <- Import.load _schemas
    resolvedExpression <- Import.load expression

    _ <- Core.throws (TypeCheck.typeOf resolvedSchemas)
    _ <- Core.throws (TypeCheck.typeOf resolvedExpression)

    let normalizedSchemas    = Normalize.normalize resolvedSchemas
    let normalizedExpression = Normalize.normalize resolvedExpression

    typeMetadata <- case decodeSchemas normalizedSchemas of
        Just typeMetadata -> return typeMetadata
        Nothing           -> Exception.throwIO NotASchemaRecord

    let schemasRewrite subExpression@(RecordLit keyValues) =
            Maybe.fromMaybe subExpression $ do
                let substitutions = Map.singleton "schemas" normalizedSchemas

                let substitutedExpression =
                        Substitution.substitute subExpression substitutions

                hash <- case TypeCheck.typeOf substitutedExpression of
                    Left _ ->
                        empty
                    Right subExpressionType ->
                        return (Import.hashExpression (Syntax.denote subExpressionType))

                (name, _default) <- Data.Map.lookup hash typeMetadata

                let diff a b | a == b    = Nothing
                             | otherwise = Just a

                let defaultedKeyValues =
                        Core.makeRecordField <$>
                        Map.fromMap (
                            Data.Map.differenceWith diff
                                (Map.toMap $ Core.recordFieldValue <$> keyValues)
                                (Map.toMap _default))

                let defaultedRecord = RecordLit defaultedKeyValues

                return (RecordCompletion (Field "schemas" $ Core.makeFieldSelection name) defaultedRecord)
        schemasRewrite subExpression =
            subExpression

    let rewrittenExpression :: Expr Src Import
        rewrittenExpression =
            fmap Void.absurd (Lens.transformOf Syntax.subExpressions schemasRewrite normalizedExpression)

    if Normalize.freeIn (V "schemas" 0) rewrittenExpression
        then return (Let (Syntax.makeBinding "schemas" _schemas) rewrittenExpression)
        else return expression

-- | Errors that can be thrown by `rewriteWithSchemas`
data SchemasError = NotASchemaRecord
    deriving (Exception)

instance Show SchemasError where
    show NotASchemaRecord =
        Util._ERROR <> ": The --schemas argument is not a record of schemas"