Simple Json encoder library wrapping `aeson`s `Encoding` in a better interface.

Enc.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}

module Json.Enc where

import Data.Aeson (Encoding, Value (..))
import Data.Aeson.Encoding qualified as AesonEnc
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap (KeyMap)
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Functor.Contravariant
import Data.Int (Int64)
import Data.Map.Strict qualified as Map
import Data.String (IsString (fromString))
import Data.Text.Lazy qualified as Lazy
import GHC.TypeLits
import PossehlAnalyticsPrelude

-- | A JSON encoder.
--
-- It is faster than going through 'Value', because 'Encoding' is just a wrapper around a @Bytes.Builder@.
-- But the @aeson@ interface for 'Encoding' is extremely bad, so let’s build a better one.
newtype Enc from = Enc (from -> Encoding)
  deriving (Num, Fractional) via (NumLiteralOnly "Enc" (Enc from))

-- | Run an 'Enc'
runEnc :: Enc from -> from -> Encoding
runEnc (Enc f) = f

-- | This is the way you can “zoom” into a data structure for encoding a subset of it.
--
-- e.g. if you have a record @Foo { fooField = Bar { barField = 42 }}@
-- you can get an @Enc Int@ with @(.fooField.barField) >$< Enc.int@ ('>$<' is an alias for `contramap`.
instance Contravariant Enc where
  contramap f (Enc e) = Enc $ \from -> e (f from)

-- | You can create an @Enc any@ that renders an 'Aeson.String' value with @OverloadedStrings@. The @any@ is unused and can take any type.
instance IsString (Enc any) where
  fromString s = constEnc (AesonEnc.string s)

-- | You can create an @Enc any@ that renders an 'Aeson.Number' value with an integer literal. The @any@ is unused and can take any type.
instance IntegerLiteral (Enc any) where
  integerLiteral i = constEnc (AesonEnc.integer i)

-- | You can create an @Enc any@ that renders an 'Aeson.Number' value with an floating point literal. The @any@ is unused and can take any type.
--
-- ATTN: Bear in mind that this will crash on repeating rationals, so only use for literals in code!
instance RationalLiteral (Enc any) where
  rationalLiteral r = constEnc (AesonEnc.scientific (r & fromRational @Scientific))

-- | Lift functions from "Data.Aeson.Encoding", or similar.
-- This is nice for smaller 'Encoding'-based constructions;
-- if you have a larger structure or you want to run a nested 'Enc',
-- you will want to use 'withEnc'.
--
-- If you need to really need to reinstate the @Enc@ with 'encode', use 'liftEnc_'.
liftEnc :: (from -> Encoding) -> Enc from
liftEnc f = Enc f

-- | Allow to embed an @Enc from@ into an 'Encoding', by passing a continuation.
--
-- For example, if you have a big object with lots of static fields, you can embed an @Enc from@ like this:
--
-- @@
-- myEnc :: Enc [Text]
-- myEnc = withEnc $ \(enc :: Enc [Text] -> Encoding) ->
--   Json.pairs [
--     … lots of fields …,
--     Json.pair
--       "someField"
--       -- here @enc@ is used
--       (enc $ Enc.list Enc.text)
--   ]
-- @@
--
-- Using the @enc@ callback to convert from the inner @Enc [Text]@
-- to the Encoding-based static construction around it.
withEnc :: ((Enc from -> Encoding) -> Encoding) -> Enc from
withEnc cont = Enc (\a -> cont (\enc -> runEnc enc a))

-- | Embed an 'Encoding' verbatim (it’s a valid JSON value)
encoding :: Enc Encoding
encoding = liftEnc id

-- | Encode a 'Value' verbatim (it’s a valid JSON value)
value :: Enc Value
value = liftEnc AesonEnc.value

-- | Encode the given constant 'Encoding'. The @any@ is unused and can take any type.
constEnc :: Encoding -> Enc any
constEnc enc = Enc $ \_any -> enc

-- | Encode an empty 'Array'. The @any@ is unused and can take any type.
emptyArray :: Enc any
emptyArray = constEnc AesonEnc.emptyArray_

-- | Encode an empty 'Object'. The @any@ is unused and can take any type.
emptyObject :: Enc any
emptyObject = constEnc AesonEnc.emptyObject_

-- | Encode a 'Text'
text :: Enc Text
text = liftEnc AesonEnc.text

-- | Encode a lazy @Text@
lazyText :: Enc Lazy.Text
lazyText = liftEnc AesonEnc.lazyText

-- | Encode a 'String'
string :: Enc String
string = liftEnc AesonEnc.string

-- | Encode as 'Null' if 'Nothing', else use the given encoder for @Just a@
orNull :: Enc a -> Enc (Maybe a)
orNull inner = liftEnc $ \case
  Nothing -> AesonEnc.null_
  Just a -> runEnc inner a

-- | Encode a list as 'Array'
list :: Enc a -> Enc [a]
list e = Enc $ \l -> AesonEnc.list (runEnc e) l

-- | Encode the given list of keys and their encoders as 'Object'.
--
-- Like with 'Map.fromList', if the list contains the same key multiple times, the last value in the list is retained:
--
-- @
-- runEnc (object [ ("foo", 42), ("foo", 23) ]) ()
-- == "{\"foo\":23}"
-- @
object :: Foldable t => t (Text, Enc from) -> Enc from
object m = Enc $ \rec ->
  AesonEnc.dict
    AesonEnc.text
    (\recEnc -> runEnc recEnc rec)
    Map.foldrWithKey
    (Map.fromList $ toList m)

-- | Construct a match for matching on a sum-type; see 'match'.
data Match where
  Match :: forall a. Enc a -> a -> Match
  deriving (Num) via (NumLiteralOnly "Match" Match)

-- | An integer literal @5 :: Match@ encodes as the json number @5@.
instance IntegerLiteral Match where
  integerLiteral = Match integer

-- |  An floating point literal @5.42 :: Match@ encodes as the json number @5.42@.
--
-- ATTN: Bear in mind that this will crash on repeating rationals, so only use for literals in code!
instance RationalLiteral Match where
  rationalLiteral r = Match (rationalLiteral @(Enc ()) r) ()

-- | An string literal @"foo" :: Match@ encodes as the json string @"foo"@.
instance IsString Match where
  fromString s = Match string s

-- | Match on a sum-type, encoding each case with the given 'Match'.
--
-- @
-- foo :: Enc (Either Text Int)
-- foo = match $ \case
--   Left t -> Match text t
--   Right i -> Match (showToText @Int >$< text) i
--
-- ex  = runEnc foo (Left "foo") == "\"foo\""
-- ex2 = runEnc foo (Right 42  ) == "\"42\""
-- @
match :: (from -> Match) -> Enc from
match f = Enc $ \from -> do
  case f from of
    Match encA a -> runEnc encA a

-- | Construct a choice match for matching on a sum-type; see 'choice'.
data Choice where
  Choice :: forall a. Text -> Enc a -> a -> Choice

-- | Encode a sum type as a @Choice@, an object with a @tag@/@value@ pair,
-- which is the conventional json sum type representation in our codebase.
--
-- @
-- foo :: Enc (Maybe Text)
-- foo = choice $ \case
--   Nothing -> Choice "no" emptyObject ()
--   Just t -> Choice "yes" text t
--
-- ex = runEnc foo Nothing == "{\"tag\": \"no\", \"value\": {}}"
-- ex2 = runEnc foo (Just "hi") == "{\"tag\": \"yes\", \"value\": \"hi\"}"
-- @
choice :: (from -> Choice) -> Enc from
choice f =
  Enc $ \from -> do
    case f from of
      Choice key encA a ->
        AesonEnc.pairs $
          mconcat
            [ AesonEnc.pair "tag" (AesonEnc.text key),
              AesonEnc.pair "value" (runEnc encA a)
            ]

-- | Encode a 'Map'.
--
-- We can’t really set the key to anything but text (We don’t keep the tag of 'Encoding')
-- so instead we allow anything that’s coercible from text as map key (i.e. newtypes).
map :: forall k v. (Coercible k Text) => Enc v -> Enc (Map k v)
map valEnc =
  Enc $ \m ->
    AesonEnc.dict
      (AesonEnc.text . coerce @k @Text)
      (runEnc valEnc)
      Map.foldrWithKey
      m

-- | Encode a 'KeyMap'
keyMap :: Enc v -> Enc (KeyMap v)
keyMap valEnc =
  Enc $ \m ->
    AesonEnc.dict
      (AesonEnc.text . Key.toText)
      (runEnc valEnc)
      KeyMap.foldrWithKey
      m

-- | Encode 'Null'. The @any@ is unused and can take any type.
null :: Enc any
null = constEnc AesonEnc.null_

-- | Encode 'Bool'. The @any@ is unused and can take any type.
bool :: Enc Bool
bool = liftEnc AesonEnc.bool

-- | Encode an 'Integer' as 'Number'.
-- TODO: is it okay to just encode an arbitrarily-sized integer into json?
integer :: Enc Integer
integer = liftEnc AesonEnc.integer

-- | Encode a 'Scientific' as 'Number'.
scientific :: Enc Scientific
scientific = liftEnc AesonEnc.scientific

-- | Encode a 'Natural' as 'Number'.
natural :: Enc Natural
natural = toInteger @Natural >$< integer

-- | Encode an 'Int' as 'Number'.
int :: Enc Int
int = liftEnc AesonEnc.int

-- | Encode an 'Int64' as 'Number'.
int64 :: Enc Int64
int64 = liftEnc AesonEnc.int64

-- | Implement this class if you want your type to only implement the part of 'Num'
-- that allows creating them from Integer-literals, then derive Num via 'NumLiteralOnly':
--
-- @
-- data Foo = Foo Integer
--   deriving (Num) via (NumLiteralOnly "Foo" Foo)
--
-- instance IntegerLiteral Foo where
--  integerLiteral i = Foo i
-- @
class IntegerLiteral a where
  integerLiteral :: Integer -> a

-- | The same as 'IntegerLiteral' but for floating point literals.
class RationalLiteral a where
  rationalLiteral :: Rational -> a

-- | Helper class for @deriving (Num) via …@, implements only literal syntax for integer and floating point numbers,
-- and throws descriptive runtime errors for any other methods in 'Num'.
--
-- See 'IntegerLiteral' and 'RationalLiteral' for examples.
newtype NumLiteralOnly (sym :: Symbol) num = NumLiteralOnly num

instance (IntegerLiteral num, KnownSymbol sym) => Num (NumLiteralOnly sym num) where
  fromInteger = NumLiteralOnly . integerLiteral
  (+) = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to add (+) (NumLiteralOnly)|]
  (*) = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to multiply (*) (NumLiteralOnly)|]
  (-) = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to subtract (-) (NumLiteralOnly)|]
  abs = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to use `abs` (NumLiteralOnly)|]
  signum = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to use `signum` (NumLiteralOnly)|]

instance (IntegerLiteral num, RationalLiteral num, KnownSymbol sym) => Fractional (NumLiteralOnly sym num) where
  fromRational = NumLiteralOnly . rationalLiteral
  recip = error [fmt|Only use as rational literal allowed for {symbolVal (Proxy @sym)}, you tried to use `recip` (NumLiteralOnly)|]
  (/) = error [fmt|Only use as numeric literal allowed for {symbolVal (Proxy @sym)}, you tried to divide (/) (NumLiteralOnly)|]