-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Facilities for re-inserting parsed comments back into an AST.
-- Actual parsing of comments is handled in "Swarm.Language.Parser.Lex".
module Swarm.Language.Parser.Comment (
  -- * Comment AST insertion
  populateComments,
  populateStandaloneComments,
  populateSuffixComments,

  -- * Generic tree traversals
  preorder,
  revpostorder,
) where

import Control.Lens (backwards, mapMOf, (%~))
import Control.Lens.Plated (Plated, plate)
import Control.Monad ((>=>))
import Control.Monad.State (MonadState (..), State, evalState)
import Data.Foldable qualified as F
import Data.List (partition)
import Data.Sequence (Seq, (<|), (|>))
import Swarm.Language.Syntax

------------------------------------------------------------
-- Comment insertion
------------------------------------------------------------

-- The approach for preserving comments is taken from
-- https://www.reddit.com/r/haskell/comments/ni4gpm/comment/gz0ipmp/ . In short:
--
--   (1) Parse all comments out-of-band and record a source span for
--       each (this is done in "Swarm.Language.Parser.Lex").
--
--   (2) For each standalone comment (i.e. comments on a line by
--       themselves), attach them to the earliest node in a preorder
--       traversal which begins after the comment.
--
--   (3) For each suffix comment (i.e. comments after something else
--       at the end of a line, or in the middle of a line), attach
--       them to the latest node in a postorder traversal which ends
--       before the comment.

-- | Re-insert parsed comments into an AST.  Prerequisite: the sequence of comments
--   must be in order by 'SrcLoc'.
populateComments :: Seq Comment -> Syntax -> Syntax
populateComments :: Seq Comment -> Syntax -> Syntax
populateComments Seq Comment
cmts = [Comment] -> Syntax -> Syntax
populateStandaloneComments [Comment]
standalone (Syntax -> Syntax) -> (Syntax -> Syntax) -> Syntax -> Syntax
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Comment] -> Syntax -> Syntax
populateSuffixComments [Comment]
suffix
 where
  ([Comment]
standalone, [Comment]
suffix) = (Comment -> Bool) -> [Comment] -> ([Comment], [Comment])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition Comment -> Bool
isStandalone (Seq Comment -> [Comment]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList Seq Comment
cmts)

-- | Insert comments from the state at the current AST node (using the
--   provided insertion function) as long as the custom comparison
--   function returns 'True' when applied to the 'SrcLoc's of the next
--   comment and the AST node (in that order).
insertComments ::
  (SrcLoc -> SrcLoc -> Bool) ->
  (Comment -> Comments -> Comments) ->
  Syntax ->
  State [Comment] Syntax
insertComments :: (SrcLoc -> SrcLoc -> Bool)
-> (Comment -> Comments -> Comments)
-> Syntax
-> State [Comment] Syntax
insertComments SrcLoc -> SrcLoc -> Bool
cmpLoc Comment -> Comments -> Comments
ins = Syntax -> State [Comment] Syntax
forall {m :: * -> *}. MonadState [Comment] m => Syntax -> m Syntax
go
 where
  go :: Syntax -> m Syntax
go s :: Syntax
s@(CSyntax SrcLoc
l Term
t Comments
cs) = do
    [Comment]
curCmts <- m [Comment]
forall s (m :: * -> *). MonadState s m => m s
get
    case [Comment]
curCmts of
      [] -> Syntax -> m Syntax
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Syntax
s
      (Comment
nextCmt : [Comment]
restCmts) -> case Comment -> SrcLoc
commentSrcLoc Comment
nextCmt SrcLoc -> SrcLoc -> Bool
`cmpLoc` SrcLoc
l of
        Bool
True -> [Comment] -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put [Comment]
restCmts m () -> m Syntax -> m Syntax
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Syntax -> m Syntax
go (SrcLoc -> Term -> Comments -> Syntax
CSyntax SrcLoc
l Term
t (Comment -> Comments -> Comments
ins Comment
nextCmt Comments
cs))
        Bool
False -> Syntax -> m Syntax
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Syntax
s

-- | Given a list of standalone comments sorted by 'SrcLoc', insert
--   them into the given AST, attaching each comment to the earliest
--   node in a preorder traversal which begins after it.
populateStandaloneComments :: [Comment] -> Syntax -> Syntax
populateStandaloneComments :: [Comment] -> Syntax -> Syntax
populateStandaloneComments [Comment]
cmts =
  (State [Comment] Syntax -> [Comment] -> Syntax)
-> [Comment] -> State [Comment] Syntax -> Syntax
forall a b c. (a -> b -> c) -> b -> a -> c
flip State [Comment] Syntax -> [Comment] -> Syntax
forall s a. State s a -> s -> a
evalState [Comment]
cmts
    (State [Comment] Syntax -> Syntax)
-> (Syntax -> State [Comment] Syntax) -> Syntax -> Syntax
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Syntax -> State [Comment] Syntax)
-> Syntax -> State [Comment] Syntax
forall a (m :: * -> *).
(Plated a, Monad m) =>
(a -> m a) -> a -> m a
preorder ((SrcLoc -> SrcLoc -> Bool)
-> (Comment -> Comments -> Comments)
-> Syntax
-> State [Comment] Syntax
insertComments SrcLoc -> SrcLoc -> Bool
srcLocStartsBefore (\Comment
c -> (Seq Comment -> Identity (Seq Comment))
-> Comments -> Identity Comments
Lens' Comments (Seq Comment)
beforeComments ((Seq Comment -> Identity (Seq Comment))
 -> Comments -> Identity Comments)
-> (Seq Comment -> Seq Comment) -> Comments -> Comments
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Seq Comment -> Comment -> Seq Comment
forall a. Seq a -> a -> Seq a
|> Comment
c)))

-- | Given a list of suffix comments sorted by 'SrcLoc', insert
--   them into the given AST, attaching each comment to the latest
--   node in a postorder traversal which ends before it.
populateSuffixComments :: [Comment] -> Syntax -> Syntax
populateSuffixComments :: [Comment] -> Syntax -> Syntax
populateSuffixComments [Comment]
cmts =
  (State [Comment] Syntax -> [Comment] -> Syntax)
-> [Comment] -> State [Comment] Syntax -> Syntax
forall a b c. (a -> b -> c) -> b -> a -> c
flip State [Comment] Syntax -> [Comment] -> Syntax
forall s a. State s a -> s -> a
evalState ([Comment] -> [Comment]
forall a. [a] -> [a]
reverse [Comment]
cmts)
    (State [Comment] Syntax -> Syntax)
-> (Syntax -> State [Comment] Syntax) -> Syntax -> Syntax
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Syntax -> State [Comment] Syntax)
-> Syntax -> State [Comment] Syntax
forall a (m :: * -> *).
(Plated a, Monad m) =>
(a -> m a) -> a -> m a
revpostorder ((SrcLoc -> SrcLoc -> Bool)
-> (Comment -> Comments -> Comments)
-> Syntax
-> State [Comment] Syntax
insertComments ((SrcLoc -> SrcLoc -> Bool) -> SrcLoc -> SrcLoc -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip SrcLoc -> SrcLoc -> Bool
srcLocEndsBefore) (\Comment
c -> (Seq Comment -> Identity (Seq Comment))
-> Comments -> Identity Comments
Lens' Comments (Seq Comment)
afterComments ((Seq Comment -> Identity (Seq Comment))
 -> Comments -> Identity Comments)
-> (Seq Comment -> Seq Comment) -> Comments -> Comments
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ (Comment
c Comment -> Seq Comment -> Seq Comment
forall a. a -> Seq a -> Seq a
<|)))

------------------------------------------------------------
-- Traversals
------------------------------------------------------------

-- $setup
-- >>> import Control.Monad.State
-- >>> import Data.Tree
-- >>> import Data.List (intercalate)
-- >>> next :: Tree Int -> State Int (Tree Int); next (Node _ cs) = do { i <- get; put (i+1); return (Node i cs) }
-- >>> showTree :: Show a => Tree a -> String; showTree = foldTree (\n cs -> show n ++ case cs of { [] -> ""; _ -> "(" ++ intercalate " " cs ++ ")" })
-- >>> exampleTree = Node 0 [Node 0 [], Node 0 [Node 0 [], Node 0 [], Node 0 []], Node 0 [Node 0 []]]

-- | Preorder traversal of a 'Plated' structure with a monadic
--   transformation.  Apply the transformation at the root, then
--   recursively transform each of the children.
--
-- >>> showTree (evalState (preorder next exampleTree) 0)
-- "0(1 2(3 4 5) 6(7))"
preorder :: (Plated a, Monad m) => (a -> m a) -> (a -> m a)
preorder :: forall a (m :: * -> *).
(Plated a, Monad m) =>
(a -> m a) -> a -> m a
preorder a -> m a
g = a -> m a
go
 where
  go :: a -> m a
go = a -> m a
g (a -> m a) -> (a -> m a) -> a -> m a
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> LensLike (WrappedMonad m) a a a a -> (a -> m a) -> a -> m a
forall (m :: * -> *) s t a b.
LensLike (WrappedMonad m) s t a b -> (a -> m b) -> s -> m t
mapMOf LensLike (WrappedMonad m) a a a a
forall a. Plated a => Traversal' a a
Traversal' a a
plate a -> m a
go

-- | Reverse postorder traversal of a 'Plated' structure with a
--   monadic transformation.  Transform the root, then apply the
--   transformation recursively to all the children in reverse order.
--
-- >>> showTree (evalState (revpostorder next exampleTree) 0)
-- "0(7 3(6 5 4) 1(2))"
revpostorder :: (Plated a, Monad m) => (a -> m a) -> (a -> m a)
revpostorder :: forall a (m :: * -> *).
(Plated a, Monad m) =>
(a -> m a) -> a -> m a
revpostorder a -> m a
g = a -> m a
go
 where
  go :: a -> m a
go = a -> m a
g (a -> m a) -> (a -> m a) -> a -> m a
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> LensLike (WrappedMonad m) a a a a -> (a -> m a) -> a -> m a
forall (m :: * -> *) s t a b.
LensLike (WrappedMonad m) s t a b -> (a -> m b) -> s -> m t
mapMOf (Optical (->) (->) (Backwards (WrappedMonad m)) a a a a
-> LensLike (WrappedMonad m) a a a a
forall (p :: * -> * -> *) (q :: * -> * -> *) (f :: * -> *) s t a b.
(Profunctor p, Profunctor q) =>
Optical p q (Backwards f) s t a b -> Optical p q f s t a b
backwards Optical (->) (->) (Backwards (WrappedMonad m)) a a a a
forall a. Plated a => Traversal' a a
Traversal' a a
plate) a -> m a
go