-- 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 :: forall (x2 :: ExpExtensionDescriptorKind)
       (x1 :: ExpExtensionDescriptorKind) 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 Traversal (XExp x1) (XExp x2) meta1 meta2
ctorL = (meta1 -> f meta2) -> Exp x1 -> f (Exp x2)
go
  where
    go :: (meta1 -> f meta2) -> Exp x1 -> f (Exp x2)
go meta1 -> f meta2
f = \case
      ExpInt XExpInt x1
x Integer
a -> meta2 -> Integer -> Exp x2
XExpInt x2 -> Integer -> Exp x2
forall (x :: ExpExtensionDescriptorKind).
XExpInt x -> Integer -> Exp x
ExpInt (meta2 -> Integer -> Exp x2) -> f meta2 -> f (Integer -> Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> meta1 -> f meta2
f meta1
XExpInt x1
x f (Integer -> Exp x2) -> f Integer -> f (Exp x2)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Integer -> f Integer
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Integer
a
      ExpString XExpString x1
x Text
a -> meta2 -> Text -> Exp x2
XExpString x2 -> Text -> Exp x2
forall (x :: ExpExtensionDescriptorKind).
XExpString x -> Text -> Exp x
ExpString (meta2 -> Text -> Exp x2) -> f meta2 -> f (Text -> Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> meta1 -> f meta2
f meta1
XExpString x1
x f (Text -> Exp x2) -> f Text -> f (Exp x2)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Text -> f Text
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
a
      ExpBytes XExpBytes x1
x ByteString
a -> meta2 -> ByteString -> Exp x2
XExpBytes x2 -> ByteString -> Exp x2
forall (x :: ExpExtensionDescriptorKind).
XExpBytes x -> ByteString -> Exp x
ExpBytes (meta2 -> ByteString -> Exp x2)
-> f meta2 -> f (ByteString -> Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> meta1 -> f meta2
f meta1
XExpBytes x1
x f (ByteString -> Exp x2) -> f ByteString -> f (Exp x2)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ByteString -> f ByteString
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
a
      ExpSeq XExpSeq x1
x [Exp x1]
a -> meta2 -> [Exp x2] -> Exp x2
XExpSeq x2 -> [Exp x2] -> Exp x2
forall (x :: ExpExtensionDescriptorKind).
XExpSeq x -> [Exp x] -> Exp x
ExpSeq (meta2 -> [Exp x2] -> Exp x2) -> f meta2 -> f ([Exp x2] -> Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> meta1 -> f meta2
f meta1
XExpSeq x1
x f ([Exp x2] -> Exp x2) -> f [Exp x2] -> f (Exp x2)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Exp x1 -> f (Exp x2)) -> [Exp x1] -> f [Exp x2]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ((meta1 -> f meta2) -> Exp x1 -> f (Exp x2)
go meta1 -> f meta2
f) [Exp x1]
a
      ExpPrim XExpPrim x1
x MichelinePrimAp x1
a ->
        meta2 -> MichelinePrimAp x2 -> Exp x2
XExpPrim x2 -> MichelinePrimAp x2 -> Exp x2
forall (x :: ExpExtensionDescriptorKind).
XExpPrim x -> MichelinePrimAp x -> Exp x
ExpPrim (meta2 -> MichelinePrimAp x2 -> Exp x2)
-> f meta2 -> f (MichelinePrimAp x2 -> Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> meta1 -> f meta2
f meta1
XExpPrim x1
x f (MichelinePrimAp x2 -> Exp x2)
-> f (MichelinePrimAp x2) -> f (Exp x2)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> LensLike
  f (MichelinePrimAp x1) (MichelinePrimAp x2) (Exp x1) (Exp x2)
-> LensLike
     f (MichelinePrimAp x1) (MichelinePrimAp x2) (Exp x1) (Exp x2)
forall (f :: * -> *) s t a b.
LensLike f s t a b -> LensLike f s t a b
L.traverseOf (([Exp x1] -> f [Exp x2])
-> MichelinePrimAp x1 -> f (MichelinePrimAp x2)
forall (x1 :: ExpExtensionDescriptorKind)
       (x2 :: ExpExtensionDescriptorKind) (f :: * -> *).
Functor f =>
([Exp x1] -> f [Exp x2])
-> MichelinePrimAp x1 -> f (MichelinePrimAp x2)
mpaArgsL (([Exp x1] -> f [Exp x2])
 -> MichelinePrimAp x1 -> f (MichelinePrimAp x2))
-> ((Exp x1 -> f (Exp x2)) -> [Exp x1] -> f [Exp x2])
-> LensLike
     f (MichelinePrimAp x1) (MichelinePrimAp x2) (Exp x1) (Exp x2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp x1 -> f (Exp x2)) -> [Exp x1] -> f [Exp x2]
forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Int (f a) (f b) a b
IndexedTraversal Int [Exp x1] [Exp x2] (Exp x1) (Exp x2)
L.traversed) ((meta1 -> f meta2) -> Exp x1 -> f (Exp x2)
go meta1 -> f meta2
f) MichelinePrimAp x1
a
      ExpX XExp x1
x -> XExp x2 -> Exp x2
forall (x :: ExpExtensionDescriptorKind). XExp x -> Exp x
ExpX (XExp x2 -> Exp x2) -> f (XExp x2) -> f (Exp x2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (meta1 -> f meta2) -> XExp x1 -> f (XExp x2)
Traversal (XExp x1) (XExp x2) meta1 meta2
ctorL meta1 -> f meta2
f XExp x1
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 :: forall meta (f :: * -> *).
Functor f =>
(meta -> f meta) -> Exp (WithMeta meta) -> f (Exp (WithMeta meta))
expMetaL meta -> f meta
f = \case
  ExpInt XExpInt (WithMeta meta)
x Integer
a ->
    meta -> f meta
f meta
XExpInt (WithMeta meta)
x f meta -> (meta -> Exp (WithMeta meta)) -> f (Exp (WithMeta meta))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \meta
x' -> XExpInt (WithMeta meta) -> Integer -> Exp (WithMeta meta)
forall (x :: ExpExtensionDescriptorKind).
XExpInt x -> Integer -> Exp x
ExpInt meta
XExpInt (WithMeta meta)
x' Integer
a
  ExpString XExpString (WithMeta meta)
x Text
a ->
    meta -> f meta
f meta
XExpString (WithMeta meta)
x f meta -> (meta -> Exp (WithMeta meta)) -> f (Exp (WithMeta meta))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \meta
x' -> XExpString (WithMeta meta) -> Text -> Exp (WithMeta meta)
forall (x :: ExpExtensionDescriptorKind).
XExpString x -> Text -> Exp x
ExpString meta
XExpString (WithMeta meta)
x' Text
a
  ExpBytes XExpBytes (WithMeta meta)
x ByteString
a ->
    meta -> f meta
f meta
XExpBytes (WithMeta meta)
x f meta -> (meta -> Exp (WithMeta meta)) -> f (Exp (WithMeta meta))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \meta
x' -> XExpBytes (WithMeta meta) -> ByteString -> Exp (WithMeta meta)
forall (x :: ExpExtensionDescriptorKind).
XExpBytes x -> ByteString -> Exp x
ExpBytes meta
XExpBytes (WithMeta meta)
x' ByteString
a
  ExpSeq XExpSeq (WithMeta meta)
x [Exp (WithMeta meta)]
a ->
    meta -> f meta
f meta
XExpSeq (WithMeta meta)
x f meta -> (meta -> Exp (WithMeta meta)) -> f (Exp (WithMeta meta))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \meta
x' -> XExpSeq (WithMeta meta)
-> [Exp (WithMeta meta)] -> Exp (WithMeta meta)
forall (x :: ExpExtensionDescriptorKind).
XExpSeq x -> [Exp x] -> Exp x
ExpSeq meta
XExpSeq (WithMeta meta)
x' [Exp (WithMeta meta)]
a
  ExpPrim XExpPrim (WithMeta meta)
x MichelinePrimAp (WithMeta meta)
a ->
    meta -> f meta
f meta
XExpPrim (WithMeta meta)
x f meta -> (meta -> Exp (WithMeta meta)) -> f (Exp (WithMeta meta))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \meta
x' -> XExpPrim (WithMeta meta)
-> MichelinePrimAp (WithMeta meta) -> Exp (WithMeta meta)
forall (x :: ExpExtensionDescriptorKind).
XExpPrim x -> MichelinePrimAp x -> Exp x
ExpPrim meta
XExpPrim (WithMeta meta)
x' MichelinePrimAp (WithMeta meta)
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 :: forall (x2 :: ExpExtensionDescriptorKind)
       (x1 :: ExpExtensionDescriptorKind) meta2 meta1.
(ExpExtrasConstrained (IsEq meta1) x1,
 ExpExtrasConstrained (IsEq meta2) x2) =>
Traversal
  (ExpressionWithMeta meta1) (ExpressionWithMeta meta2) meta1 meta2
expAllMetaL = Traversal
  (XExp (WithMeta meta1)) (XExp (WithMeta meta2)) meta1 meta2
-> Traversal
     (Exp (WithMeta meta1)) (Exp (WithMeta meta2)) meta1 meta2
forall (x2 :: ExpExtensionDescriptorKind)
       (x1 :: ExpExtensionDescriptorKind) 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 Over (->) f Void Void meta1 meta2
(meta1 -> f meta2)
-> XExp (WithMeta meta1) -> f (XExp (WithMeta meta2))
forall {k} (p :: k -> * -> *) (f :: * -> *) (a :: k) b.
Over p f Void Void a b
Traversal
  (XExp (WithMeta meta1)) (XExp (WithMeta meta2)) meta1 meta2
L.devoid

-- | Lift plain 'Expression' to 'ExpressionWithMeta'.
expAnnotate :: Expression -> ExpressionWithMeta ()
expAnnotate :: Expression -> ExpressionWithMeta ()
expAnnotate = Traversal (XExp RegularExp) (XExp (WithMeta ())) () ()
-> Traversal Expression (ExpressionWithMeta ()) () ()
forall (x2 :: ExpExtensionDescriptorKind)
       (x1 :: ExpExtensionDescriptorKind) 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 Over (->) f Void Void () ()
(() -> f ()) -> XExp RegularExp -> f (XExp (WithMeta ()))
forall {k} (p :: k -> * -> *) (f :: * -> *) (a :: k) b.
Over p f Void Void a b
Traversal (XExp RegularExp) (XExp (WithMeta ())) () ()
L.devoid ((() -> Identity ())
 -> Expression -> Identity (ExpressionWithMeta ()))
-> (() -> ()) -> Expression -> ExpressionWithMeta ()
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ () -> ()
forall a. a -> a
id