Annotation.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# LANGUAGE UndecidableSuperClasses #-}

-- | Type and field annotations for Lorentz types.
module Lorentz.Annotation
  ( AnnOptions (..)
  , dropPrefix
  , appendTo
  , toCamel
  , toPascal
  , toSnake
  , stripFieldPrefix

  , ctorNameToAnnWithOptions
  , FollowEntrypointFlag (..)
  , GenerateFieldAnnFlag (..)

  , HasAnnotation (..)
  , GHasAnnotation (..)
  , gGetAnnotationNoField
  , insertTypeAnn
  ) where

import Data.Default (Default(..))
import GHC.Generics qualified as G

import Morley.Michelson.Text
import Morley.Michelson.Typed
  (BigMap, BigMapId, ContractRef(..), EpAddress, KnownIsoT, Notes(..), Operation, Ticket, ToT,
  insertTypeAnn, starNotes)
import Morley.Michelson.Typed.Haskell.Doc (TypeHasFieldNamingStrategy(..))
import Morley.Michelson.Typed.Haskell.Value (GValueType)
import Morley.Michelson.Untyped (FieldAnn, TypeAnn, VarAnn, mkAnnotation, noAnn)
import Morley.Tezos.Address
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Util.Generic
import Morley.Util.Named
import Morley.Util.Text
import Morley.Util.TypeLits

{- $setup
>>> import Lorentz
-}

----------------------------------------------------------------------------
-- Annotation Customization
----------------------------------------------------------------------------

-- | Allow customization of field annotation generated for a type
-- when declaring its 'HasAnnotation' instance.
data AnnOptions = AnnOptions
  { fieldAnnModifier :: Text -> Text
  }

instance Default AnnOptions where
  def = AnnOptions id

-- | @appendTo suffix fields field@ appends the given suffix to @field@
-- if the field exists in the @fields@ list.
appendTo :: Text -> [Text] -> Text -> Text
appendTo suffix fields field
  | field `elem` fields = field <> suffix
  | otherwise = field

----------------------------------------------------------------------------
-- Typeclasses related to Annotaiton Generation
----------------------------------------------------------------------------

ctorNameToAnnWithOptions :: forall ctor. (KnownSymbol ctor, HasCallStack) => AnnOptions -> FieldAnn
ctorNameToAnnWithOptions o = unsafe . mkAnnotation . fieldAnnModifier o $ headToLower $ (symbolValT' @ctor)

-- | Used in `GHasAnnotation` and `HasAnnotation` as a flag to track
-- whether or not it directly follows an entrypoint to avoid introducing
-- extra entrypoints.
data FollowEntrypointFlag = FollowEntrypoint | NotFollowEntrypoint

-- | Used in `GHasAnnotation` as a flag to track whether or not field/constructor
-- annotations should be generated.
data GenerateFieldAnnFlag = GenerateFieldAnn | NotGenerateFieldAnn

-- | Use this in the instance of @HasAnnotation@ when field annotations
-- should not be generated.
gGetAnnotationNoField
    :: forall a. (GHasAnnotation (GRep a), GValueType (GRep a) ~ ToT a)
    => FollowEntrypointFlag -> Notes (ToT a)
gGetAnnotationNoField =
  \_ -> gGetAnnotation @(GRep a) def NotFollowEntrypoint NotGenerateFieldAnn ^. _1

{- | This class defines the type and field annotations for a given type. Right
now the type annotations come from names in a named field, and field annotations
are generated from the record fields.

Allows generic derivation. The type must have @Generic@ and @IsoValue@
instances.

>>> data Foo = Foo
>>> instance HasAnnotation Foo
...
... GHC.Generics.Rep Foo
... is stuck. Likely
... Generic Foo
... instance is missing or out of scope.
...

>>> data Foo = Foo deriving (Generic, IsoValue)
>>> instance HasAnnotation Foo
-}
class HasAnnotation a where
  getAnnotation :: FollowEntrypointFlag -> Notes (ToT a)
  default getAnnotation
    :: ( GHasAnnotation (GRep a), GValueType (GRep a) ~ ToT a
       , GIsSumType (GRep a), TypeHasFieldNamingStrategy a
       )
    => FollowEntrypointFlag
    -> Notes (ToT a)
  getAnnotation b =
    gGetAnnotation @(GRep a) (fromMaybe (gAnnOptions @a) $ annOptions @a) b GenerateFieldAnn ^. _1

  annOptions :: Maybe AnnOptions
  annOptions = Nothing

gAnnOptions :: forall a. (GIsSumType (GRep a), TypeHasFieldNamingStrategy a) => AnnOptions
gAnnOptions
  | gIsSumType @(GRep a) = def
  | otherwise = def { fieldAnnModifier = typeFieldNamingStrategy @a }

instance (HasAnnotation a, KnownSymbol name)
  => HasAnnotation (NamedF Identity a name) where
  getAnnotation b = insertTypeAnn (symbolAnn @name) $
    getAnnotation @a b
    where
      symbolAnn :: forall s. KnownSymbol s => TypeAnn
      symbolAnn = unsafe . mkAnnotation $ symbolValT' @s

instance (HasAnnotation (Maybe a), KnownSymbol name)
  => HasAnnotation (NamedF Maybe a name) where
  getAnnotation b = getAnnotation @(NamedF Identity (Maybe a) name) b

-- Primitive instances
instance (HasAnnotation a) => HasAnnotation (Maybe a) where
  getAnnotation _ = NTOption noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation ()

instance HasAnnotation Integer where
  getAnnotation _ = starNotes

instance HasAnnotation Natural where
  getAnnotation _ = starNotes

instance HasAnnotation MText where
  getAnnotation _ = starNotes

instance HasAnnotation Bool where
  getAnnotation _ = starNotes

instance HasAnnotation ByteString where
  getAnnotation _ = starNotes

instance HasAnnotation Mutez where
  getAnnotation _ = starNotes

instance HasAnnotation Address where
  getAnnotation _ = starNotes

instance HasAnnotation EpAddress where
  getAnnotation _ = starNotes

instance HasAnnotation KeyHash where
  getAnnotation _ = starNotes

instance HasAnnotation Timestamp where
  getAnnotation _ = starNotes

instance HasAnnotation PublicKey where
  getAnnotation _ = starNotes

instance HasAnnotation Signature where
  getAnnotation _ = starNotes

instance HasAnnotation ChainId where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation (ContractRef a) where
  getAnnotation _ = NTContract noAnn (getAnnotation @a NotFollowEntrypoint)

instance (HasAnnotation d) => HasAnnotation (Ticket d) where
  getAnnotation _ = NTTicket noAnn (getAnnotation @d NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (Map k v) where
  getAnnotation _ = NTMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (BigMap k v) where
  getAnnotation _ = NTBigMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance HasAnnotation (BigMapId k v) where
  getAnnotation _ = starNotes

instance (KnownIsoT v) => HasAnnotation (Set v) where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation [a] where
  getAnnotation _ = NTList noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation Operation where
  getAnnotation _ = starNotes

instance HasAnnotation Chest where
  getAnnotation _ = starNotes

instance HasAnnotation ChestKey where
  getAnnotation _ = starNotes

instance (HasAnnotation a, HasAnnotation b) => HasAnnotation (Either a b)

instance (HasAnnotation a, HasAnnotation b) => HasAnnotation (a, b)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c) => HasAnnotation (a, b, c)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d) => HasAnnotation (a, b, c, d)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e)
  => HasAnnotation (a, b, c, d, e)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e, HasAnnotation f)
  => HasAnnotation (a, b, c, d, e, f)
instance
  ( HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e
  , HasAnnotation f, HasAnnotation g)
  => HasAnnotation (a, b, c, d, e, f, g)

-- | A Generic @HasAnnotation@ implementation
class GHasAnnotation a where
  gGetAnnotation :: AnnOptions -> FollowEntrypointFlag -> GenerateFieldAnnFlag -> (Notes (GValueType a), FieldAnn, VarAnn)

instance GHasAnnotation G.U1 where
  gGetAnnotation _ _ _ = (starNotes, noAnn, noAnn)

instance (GHasAnnotation x)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel 'Nothing b c d) x)
  where
  gGetAnnotation o b b2 = gGetAnnotation @x o b b2

instance (GHasAnnotation x, KnownSymbol a)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel ('Just a) b c d) x)
  where
  gGetAnnotation o b b2 = case b2 of
    GenerateFieldAnn -> (gGetAnnotation @x o b b2 ^. _1, ctorNameToAnnWithOptions @a o, noAnn)
    NotGenerateFieldAnn -> (gGetAnnotation @x o b b2 ^. _1, noAnn, noAnn)

instance (GHasAnnotation x, KnownSymbol a) => GHasAnnotation (G.M1 G.C ('G.MetaCons a _p _f) x) where
  gGetAnnotation o b b2 =
    ( gGetAnnotation @x o b b2 ^. _1
    , ctorNameToAnnWithOptions @a o
    , noAnn
    )

instance (GHasAnnotation x) => GHasAnnotation (G.M1 G.D i1 x) where
  gGetAnnotation o b b2 = gGetAnnotation @x o b b2

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:+: y)
  where

  gGetAnnotation o followEntrypointFlag generateAnnFlag =
    let (xTypeAnn, xFieldAnn, xVarAnn) = gGetAnnotation @x o followEntrypointFlag generateAnnFlag
        (yTypeAnn, yFieldAnn, yVarAnn) = gGetAnnotation @y o followEntrypointFlag generateAnnFlag
    in case (followEntrypointFlag, generateAnnFlag) of
          (NotFollowEntrypoint, GenerateFieldAnn) ->
            ( NTOr noAnn
                xFieldAnn yFieldAnn
                xTypeAnn yTypeAnn
            , noAnn
            , xVarAnn
            )
          _ ->
            ( NTOr noAnn
                noAnn noAnn
                xTypeAnn yTypeAnn
            , noAnn
            , yVarAnn
            )

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:*: y)
  where
  gGetAnnotation o _ b2 =
    let  (xTypeAnn, xFieldAnn, xVarAnn) = gGetAnnotation @x o NotFollowEntrypoint b2
         (yTypeAnn, yFieldAnn, yVarAnn) = gGetAnnotation @y o NotFollowEntrypoint b2
    in
      ( NTPair noAnn
          xFieldAnn yFieldAnn
          xVarAnn yVarAnn
          xTypeAnn yTypeAnn
      , noAnn
      , noAnn
      )

instance (HasAnnotation x) => GHasAnnotation (G.Rec0 x) where
  gGetAnnotation _ b _ = (getAnnotation @x b, noAnn, noAnn)

-- | Check if a 'Generic' type is a sum type.
class GIsSumType a where
  gIsSumType :: Bool

instance GIsSumType G.U1 where gIsSumType = False
instance GIsSumType G.V1 where gIsSumType = False
instance GIsSumType (x G.:+: y) where gIsSumType = True
instance GIsSumType (x G.:*: y) where gIsSumType = False
instance GIsSumType (G.S1 _a x) where gIsSumType = False
instance GIsSumType (G.C1 _a x) where gIsSumType = False
instance GIsSumType (G.Rec0 x) where gIsSumType = False
instance GIsSumType x => GIsSumType (G.D1 _a x) where gIsSumType = gIsSumType @x