comment.scala

object Test extends App {
  import scala.meta._
  import scala.meta.tokens.Token.Comment
  val tree = """
    |package io.buildo.baseexample
    |
    |package enums
    |
    |import ingredients.caseenum.annotations.enum
    |
    |// scalastyle:off number.of.types
    |@enum trait CampingLocation {
    |  // Foobar
    |  object Seaside // trailing
    |  object Mountains
    |}
  """.stripMargin.parse[Source].get

  case class AssociatedComments(leadingMap: Map[Token, Seq[Comment]],
                                trailingMap: Map[Token, Seq[Comment]]) {
    def leading(tree: Tree): Seq[Comment] =
      (for {
        token <- tree.tokens.headOption
        comments <- leadingMap.get(token)
      } yield comments).getOrElse(Seq.empty[Comment])

    def trailing(tree: Tree): Seq[Comment] =
      (for {
        token <- tree.tokens.lastOption
        comments <- trailingMap.get(token)
      } yield comments).getOrElse(Seq.empty[Comment])

    def contains(tree: Tree) =
      trailing(tree).nonEmpty || leading(tree).nonEmpty
  }

  def getAssociatedComments(tokens: Tokens): AssociatedComments = {
    import scala.meta.tokens.Token._
    val leadingBuilder = Map.newBuilder[Token, Seq[Comment]]
    val trailingBuilder = Map.newBuilder[Token, Seq[Comment]]
    val leading = Seq.newBuilder[Comment]
    val trailing = Seq.newBuilder[Comment]
    var isLeading = true
    var lastToken: Token = tokens.head
    tokens.foreach {
      case c: Comment =>
        if (isLeading) leading += c
        else trailing += c
      case _: `\n` => isLeading = true
      case _: Trivia =>
      case currentToken =>
        val t = trailing.result()
        if (t.nonEmpty) {
          trailingBuilder += lastToken -> trailing.result()
          trailing.clear()
        }
        val l = leading.result()
        if (l.nonEmpty) {
          leadingBuilder += currentToken -> leading.result()
          leading.clear()
        }
        lastToken = currentToken
        isLeading = false
    }
    AssociatedComments(leadingBuilder.result(), trailingBuilder.result())
  }

  val comments = getAssociatedComments(tree.tokens)
  tree.transform {
    case t: Tree if comments.contains(t) =>
      println(t, t.getClass, comments.leading(t), comments.trailing(t))
      t
  }
  println(comments.leadingMap)
  println(comments.trailingMap)
}

Output:

[debug] CommentTest.scala:99      [t]: @enum trait CampingLocation {
  // Foobar
  object Seaside // trailing
  object Mountains
}
[debug] CommentTest.scala:99      [t.getClass]: class scala.meta.Defn$Trait$Impl
[debug] CommentTest.scala:99      [comments.leading(t)]: List(// scalastyle:off number.of.types (93..126))
[debug] CommentTest.scala:99      [comments.trailing(t)]: List()
[debug] CommentTest.scala:99      [t]: @enum
[debug] CommentTest.scala:99      [t.getClass]: class scala.meta.Mod$Annot$Impl
[debug] CommentTest.scala:99      [comments.leading(t)]: List(// scalastyle:off number.of.types (93..126))
[debug] CommentTest.scala:99      [comments.trailing(t)]: List()
[debug] CommentTest.scala:99      [t]: object Seaside
[debug] CommentTest.scala:99      [t.getClass]: class scala.meta.Defn$Object$Impl
[debug] CommentTest.scala:99      [comments.leading(t)]: List(// Foobar (159..168))
[debug] CommentTest.scala:99      [comments.trailing(t)]: List(// trailing (186..197))
[debug] CommentTest.scala:99      [t]: Seaside
[debug] CommentTest.scala:99      [t.getClass]: class scala.meta.Term$Name$Impl
[debug] CommentTest.scala:99      [comments.leading(t)]: List()
[debug] CommentTest.scala:99      [comments.trailing(t)]: List(// trailing (186..197))
[debug] CommentTest.scala:102     [comments.leadingMap]: Map(@ (127..128) -> List(// scalastyle:off number.of.types (93..126)), object (171..177) -> List(// Foobar (159..168)))
[debug] CommentTest.scala:103     [comments.trailingMap]: Map(Seaside (178..185) -> List(// trailing (186..197)))