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

-- | Expression that carries the same extra data for all 'Exp' constructors.
module Morley.Micheline.Expression.WithMeta
  ( -- * General case
    expAllExtraL

    -- * Custom extra field, no extra constructors
  , WithMeta
  , ExpressionWithMeta
  , expMetaL
  , expAnnotate
  , expAllMetaL

    -- * Utilities
  , IsEq
  ) where

import Control.Lens qualified as L

import Morley.Micheline.Expression
import Morley.Util.Type (IsEq)

--------------------------------------------------------------------------------
-- General case
--------------------------------------------------------------------------------

-- | Traversal that visits all the extra fields ('XExpInt' and others)
-- in DFS order assuming they are the same for every constructor.
--
-- It is generic enough to work not only with 'ExpressionWithMeta', but
-- with any 'Exp' that has all the extra fields of the same type;
-- hence, this traversal is applicable to 'Expression' too.
--
-- This also supports additional constructors. Use 'L.devoid' if you have none.
--
-- Moreover, when used as setter, it can change the type of meta.
expAllExtraL
  :: forall x2 x1 meta2 meta1.
     ( ExpExtrasConstrained (IsEq meta1) x1
     , ExpExtrasConstrained (IsEq meta2) x2
     )
  => Traversal (XExp x1) (XExp x2) meta1 meta2
  -> Traversal (Exp x1) (Exp x2) meta1 meta2
expAllExtraL ctorL = go
  where
    go f = \case
      ExpInt x a -> ExpInt <$> f x <*> pure a
      ExpString x a -> ExpString <$> f x <*> pure a
      ExpBytes x a -> ExpBytes <$> f x <*> pure a
      ExpSeq x a -> ExpSeq <$> f x <*> traverse (go f) a
      ExpPrim x a ->
        ExpPrim <$> f x <*> L.traverseOf (mpaArgsL . L.traversed) (go f) a
      ExpX x -> ExpX <$> ctorL f x

--------------------------------------------------------------------------------
-- Custom extra field, no extra constructors
--------------------------------------------------------------------------------

-- | Expression that has the same type of metadata attached to each of
-- its constructors.
data WithMeta (meta :: Type) :: ExpExtensionDescriptorKind
instance ExpExtensionDescriptor (WithMeta m) where
  type XExpInt (WithMeta m) = m
  type XExpString (WithMeta m) = m
  type XExpBytes (WithMeta m) = m
  type XExpSeq (WithMeta m) = m
  type XExpPrim (WithMeta m) = m

-- | Alias for expression with given meta.
type ExpressionWithMeta meta = Exp (WithMeta meta)

{-# ANN module ("HLint: ignore Avoid lambda using `infix`" :: Text) #-}

-- | Lens for getting immediate meta of the node.
expMetaL :: Lens' (Exp (WithMeta meta)) meta
expMetaL f = \case
  ExpInt x a ->
    f x <&> \x' -> ExpInt x' a
  ExpString x a ->
    f x <&> \x' -> ExpString x' a
  ExpBytes x a ->
    f x <&> \x' -> ExpBytes x' a
  ExpSeq x a ->
    f x <&> \x' -> ExpSeq x' a
  ExpPrim x a ->
    f x <&> \x' -> ExpPrim x' a

-- | Traversal that visits all the metas in DFS order.
-- This is a specialization of 'expAllExtraL'.
--
-- This is pretty similar to 'expMetaL', but picks meta of all the
-- transitive children.
--
-- When used as setter, it can change the type of meta.
expAllMetaL
  :: forall x2 x1 meta2 meta1.
     ( ExpExtrasConstrained (IsEq meta1) x1
     , ExpExtrasConstrained (IsEq meta2) x2
     )
  => Traversal (ExpressionWithMeta meta1) (ExpressionWithMeta meta2) meta1 meta2
expAllMetaL = expAllExtraL L.devoid

-- | Lift plain 'Expression' to 'ExpressionWithMeta'.
expAnnotate :: Expression -> ExpressionWithMeta ()
expAnnotate = expAllExtraL L.devoid %~ id