Origin: add {text,tokens}Opt methods (#3653)

In reality, they are non-trivial only for Origin.Parsed.
scalameta · Mar 16, 2024 · 1c1dd4f · 1c1dd4f
1 parent eeb1ea6
commit 1c1dd4f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 25 deletions.
diff --git a/scalameta/trees/shared/src/main/scala/scala/meta/internal/prettyprinters/TreeSyntax.scala b/scalameta/trees/shared/src/main/scala/scala/meta/internal/prettyprinters/TreeSyntax.scala
@@ -665,20 +665,16 @@ object TreeSyntax {
         )
       case t: Type.AnonymousLambda => s(t.tpe)
       case t: Type.AnonymousParam =>
-        val useStar = dialect.allowStarAsTypePlaceholder && (t.origin match {
-          case o: Origin.Parsed => !o.tokens.lastOption.exists(_.is[Token.Underscore])
-          case _ => false
-        })
+        val useStar = dialect.allowStarAsTypePlaceholder &&
+          t.origin.tokensOpt.exists(!_.lastOption.exists(_.is[Token.Underscore]))
         val ph = if (useStar) "*" else "_"
         m(SimpleTyp, o(t.variant), ph)
       case t: Type.Wildcard =>
         /* In order not to break existing tools `.syntax` should still return
          * `_` instead `?` unless specifically used.
          */
-        def questionMarkUsed = t.origin match {
-          case o: Origin.Parsed => !o.tokens.headOption.exists(_.is[Token.Underscore])
-          case _ => false
-        }
+        def questionMarkUsed =
+          t.origin.tokensOpt.exists(!_.headOption.exists(_.is[Token.Underscore]))
         val useQM = dialect.allowQuestionMarkAsTypeWildcard &&
           (dialect.allowUnderscoreAsTypePlaceholder || questionMarkUsed)
         m(SimpleTyp, s(kw(if (useQM) "?" else "_")), t.bounds)
@@ -688,10 +684,7 @@ object TreeSyntax {
         /* In order not to break existing tools `.syntax` should still return
          * `_` instead `?` unless specifically used.
          */
-        def questionMarkUsed = t.origin match {
-          case o: Origin.Parsed => !o.tokens.exists(_.is[Token.Underscore])
-          case _ => false
-        }
+        def questionMarkUsed = t.origin.tokensOpt.exists(!_.exists(_.is[Token.Underscore]))
         val useQM = dialect.allowQuestionMarkAsTypeWildcard &&
           (dialect.allowUnderscoreAsTypePlaceholder || questionMarkUsed)
         m(SimpleTyp, s(kw(if (useQM) "?" else "_"), t.bounds))
@@ -1148,7 +1141,7 @@ object TreeSyntax {
     // I expect to improve on this in the nearest future, because we had it much better until recently.
     Syntax { (x: T) =>
       x.origin match {
-        case o: Origin.Parsed if o.dialect.isEquivalentTo(dialect) => s(o.position.text)
+        case o: Origin.Parsed if o.dialect.isEquivalentTo(dialect) => s(o.text)
         case _ => reprint(x)(dialect)
       }
     }

diff --git a/scalameta/trees/shared/src/main/scala/scala/meta/internal/trees/InternalTrees.scala b/scalameta/trees/shared/src/main/scala/scala/meta/internal/trees/InternalTrees.scala
@@ -75,10 +75,7 @@ trait InternalTree extends Product {
     tokenCache.getOrElseUpdate(dialect, tokenizeForDialect(dialect))
 
   private lazy val tokensOpt: Option[Tokens] =
-    origin match {
-      case x: Origin.Parsed => Some(x.tokens)
-      case _ => origin.dialectOpt.map(tokenizeForDialect)
-    }
+    origin.tokensOpt.orElse(origin.dialectOpt.map(tokenizeForDialect))
 
   private def tokenizeForDialect(dialect: Dialect): Tokens =
     this match {
@@ -108,10 +105,7 @@ trait InternalTree extends Product {
     TreeSyntax.reprint(this)(dialect).toString
 
   private lazy val textOpt: Option[String] =
-    origin match {
-      case x: Origin.Parsed => Some(x.position.text)
-      case _ => origin.dialectOpt.map(reprintSyntax)
-    }
+    origin.textOpt.orElse(origin.dialectOpt.map(reprintSyntax))
 
   // ==============================================================
   // Intellij-friendly stubs.

diff --git a/scalameta/trees/shared/src/main/scala/scala/meta/trees/Origin.scala b/scalameta/trees/shared/src/main/scala/scala/meta/trees/Origin.scala
@@ -12,13 +12,17 @@ import scala.meta.tokens._
 trait Origin extends Optional {
   def position: Position
   def dialectOpt: Option[Dialect]
+  private[meta] def textOpt: Option[String]
+  private[meta] def tokensOpt: Option[Tokens]
 }
 
 object Origin {
   @adt.none
   object None extends Origin {
     val position: Position = Position.None
     val dialectOpt: Option[Dialect] = scala.None
+    private[meta] val textOpt: Option[String] = scala.None
+    private[meta] val tokensOpt: Option[Tokens] = scala.None
   }
 
   // `begTokenIdx` and `endTokenIdx` are half-open interval of index range
@@ -34,13 +38,13 @@ object Origin {
     }
 
     def dialectOpt: Option[Dialect] = Some(dialect)
-
-    def tokens: Tokens = {
-      allInputTokens().slice(begTokenIdx, endTokenIdx)
-    }
+    private[meta] def textOpt: Option[String] = Some(text)
+    private[meta] def tokensOpt: Option[Tokens] = Some(tokens)
 
     @inline def input: Input = source.input
     @inline def dialect: Dialect = source.dialect
+    @inline def text: String = position.text
+    def tokens: Tokens = allInputTokens().slice(begTokenIdx, endTokenIdx)
   }
 
   class ParsedSource(val input: Input)(implicit val dialect: Dialect) {
@@ -52,6 +56,8 @@ object Origin {
   class DialectOnly(dialect: Dialect) extends Origin {
     val position: Position = Position.None
     def dialectOpt: Option[Dialect] = Some(dialect)
+    private[meta] val textOpt: Option[String] = scala.None
+    private[meta] val tokensOpt: Option[Tokens] = scala.None
   }
 
   object DialectOnly {