<commit_before>package dotty.tools
package dotc
package printing

import dotty.tools.dotc.core.Contexts.Context
import dotty.tools.dotc.core.StdNames._
import dotty.tools.dotc.parsing.Parsers.Parser
import dotty.tools.dotc.parsing.Scanners.Scanner
import dotty.tools.dotc.parsing.Tokens._
import dotty.tools.dotc.reporting.Reporter
import dotty.tools.dotc.reporting.diagnostic.MessageContainer
import dotty.tools.dotc.util.Positions.Position

import scala.annotation.switch
import scala.collection.mutable.StringBuilder
import util.{Chars, SourceFile}

import scala.collection.mutable

/** This object provides functions for syntax highlighting in the REPL */
object SyntaxHighlighting {

  // Keep in sync with SyntaxHighlightingTests
  val NoColor         = Console.RESET
  val CommentColor    = Console.BLUE
  val KeywordColor    = Console.YELLOW
  val ValDefColor     = Console.CYAN
  val LiteralColor    = Console.RED
  val StringColor     = Console.GREEN
  val TypeColor       = Console.MAGENTA
  val AnnotationColor = Console.MAGENTA

  private def none(str: String) = str
  private def keyword(str: String) = KeywordColor + str + NoColor
  private def typeDef(str: String) = TypeColor + str + NoColor
  private def literal(str: String) = LiteralColor + str + NoColor
  private def valDef(str: String) = ValDefColor + str + NoColor
  private def operator(str: String) = TypeColor + str + NoColor
  private def annotation(str: String) =
    if (str.trim == "@") str else { AnnotationColor + str + NoColor }
  private val tripleQs = Console.RED_B + "???" + NoColor

  private val keywords: Seq[String] = for {
    index <- IF to ERASED // All alpha keywords
  } yield tokenString(index)

  private val interpolationPrefixes =
    'A' :: 'B' :: 'C' :: 'D' :: 'E' :: 'F' :: 'G' :: 'H' :: 'I' :: 'J' :: 'K' ::
    'L' :: 'M' :: 'N' :: 'O' :: 'P' :: 'Q' :: 'R' :: 'S' :: 'T' :: 'U' :: 'V' ::
    'W' :: 'X' :: 'Y' :: 'Z' :: '$' :: '_' :: 'a' :: 'b' :: 'c' :: 'd' :: 'e' ::
    'f' :: 'g' :: 'h' :: 'i' :: 'j' :: 'k' :: 'l' :: 'm' :: 'n' :: 'o' :: 'p' ::
    'q' :: 'r' :: 's' :: 't' :: 'u' :: 'v' :: 'w' :: 'x' :: 'y' :: 'z' :: Nil

  private val typeEnders =
   '{'  :: '}' :: ')' :: '(' :: '[' :: ']' :: '=' :: ' ' :: ',' :: '.' :: '|' ::
   '&' :: '\n' :: Nil

  def apply(chars: Iterable[Char])(implicit ctx: Context): Iterable[Char] = {
    if (ctx.settings.color.value != "never") highlight(chars)
    else chars
  }

  def highlight(chars: Iterable[Char]): Iterable[Char] = {
    var prev: Char = 0
    var remaining  = chars.toStream
    val newBuf     = new StringBuilder
    var lastValDefToken  = ""

    @forceInline def keywordStart =
      prev == 0    || prev == ' ' || prev == '{' || prev == '(' ||
      prev == '\n' || prev == '[' || prev == ',' || prev == ':' ||
      prev == '|'  || prev == '&' || prev.isDigit

    @forceInline def numberStart(c: Char) =
      c.isDigit && (!prev.isLetter || prev == '.' || prev == ' ' || prev == '(' || prev == '\u0000')

    def takeChar(): Char = takeChars(1).head
    def takeChars(x: Int): Seq[Char] = {
      val taken = remaining.take(x)
      remaining = remaining.drop(x)
      taken
    }

    while (remaining.nonEmpty) {
      val n = takeChar()
      if (interpolationPrefixes.contains(n)) {
        // Interpolation prefixes are a superset of the keyword start chars
        val (prefix, after) = remaining.span(interpolationPrefixes.contains)
        if (after.startsWith("\"")) {
          newBuf += n ++= prefix
          prev = prefix.lastOption.getOrElse(n)
          if (remaining.nonEmpty) takeChars(prefix.length + 1) // drop 1 for appendLiteral
          appendString('"', after.startsWith("\"\"\""), true)
        } else {
          if (n.isUpper && (keywordStart || prev == '.')) {
            appendWhile(n, !typeEnders.contains(_), typeDef)
          } else if (keywordStart) {
            append(n, keywords.contains(_), { kw =>
              if (kw == "new") typeDef(kw) else keyword(kw)
            })
          } else {
            newBuf += n
            prev = n
          }
        }
      } else {
        (n: @switch) match {
          case '/'  =>
            if (remaining.nonEmpty) {
              remaining.head match {
                case '/' =>
                  takeChar()
                  eolComment()
                case '*' =>
                  takeChar()
                  blockComment()
                case x =>
                  newBuf += '/'
              }
            } else newBuf += '/'
          case '='  =>
            append('=', _ == "=>", operator)
          case '<'  =>
            append('<', { x => x == "<-" || x == "<:" || x == "<%" }, operator)
          case '>'  =>
            append('>', { x => x == ">:" }, operator)
          case '#'  =>
            if (prev != ' ' && prev != '.') newBuf append operator("#")
            else newBuf += n
            prev = '#'
          case '@'  =>
            appendWhile('@', !typeEnders.contains(_), annotation)
          case '\"' =>
            appendString('\"', multiline = remaining.take(2).mkString == "\"\"", false)
          case '\'' =>
            appendSingleQuote('\'')
          case '`'  =>
            appendTo('`', _ == '`', none)
          case _    => {
            if (n == '?' && remaining.take(2).mkString == "??") {
              takeChars(2)
              newBuf append tripleQs
              prev = '?'
            }
            else if (n.isUpper && keywordStart)
              appendWhile(n, !typeEnders.contains(_), typeDef)
            else if (numberStart(n)) {
              def isNumber(c: Char): Boolean =
                c.isDigit || c == '\u0000' || (c == '.' && remaining.nonEmpty && remaining.head.isDigit)
              appendWhile(n, isNumber, literal)
            } else
              newBuf += n; prev = n
          }
        }
      }
    }

    def eolComment() = {
      newBuf append (CommentColor + "//")
      var curr = '/'
      while (curr != '\n' && remaining.nonEmpty) {
        curr = takeChar()
        newBuf += curr
      }
      prev = curr
      newBuf append NoColor
    }

    def blockComment() = {
      newBuf append (CommentColor + "/*")
      var curr = '*'
      var open = 1
      while (open > 0 && remaining.nonEmpty) {
        curr = takeChar()
        if (curr == '@') {
          appendWhile('@', !typeEnders.contains(_), annotation)
          newBuf append CommentColor
        }
        else newBuf += curr

        if (curr == '*' && remaining.nonEmpty) {
          curr = takeChar()
          newBuf += curr
          if (curr == '/') open -= 1
        } else if (curr == '/' && remaining.nonEmpty) {
          curr = takeChar()
          newBuf += curr
          if (curr == '*') open += 1
        }

        if (Chars.isLineBreakChar(curr)) {
          newBuf append CommentColor
        }
      }
      prev = curr
      newBuf append NoColor
    }

    def appendString(delim: Char, multiline: Boolean = false, inInterpolation: Boolean) = {
      var curr: Char      = 0
      var continue        = true
      var closing         = 0
      newBuf append (LiteralColor + delim)

      def shouldInterpolate =
        inInterpolation && curr == '$' && prev != '$' && remaining.nonEmpty

      def interpolate() = {
        val next = takeChar()
        if (next == '$') {
          newBuf += curr
          newBuf += next
          prev = '$'
        } else if (next == '{') {
          var open = 1 // keep track of open blocks
          newBuf append (ValDefColor + curr)
          newBuf += next
          while (remaining.nonEmpty && open > 0) {
            var c = takeChar()
            newBuf += c
            if (c == '}') open -= 1
            else if (c == '{') open += 1
          }
          newBuf append LiteralColor
        } else {
          newBuf append (ValDefColor + curr)
          newBuf += next
          var c: Char = 'a'
          while (c.isLetterOrDigit && remaining.nonEmpty) {
            c = takeChar()
            if (c != '"') newBuf += c
          }
          newBuf append LiteralColor
          if (c == '"') {
            newBuf += c
            continue = false
          }
        }
        closing = 0
      }

      while (continue && remaining.nonEmpty) {
        curr = takeChar()
        if (curr == '\\' && remaining.nonEmpty) {
          val next = takeChar()
          newBuf append (KeywordColor + curr)
          if (next == 'u') {
            val code = "u" + takeChars(4).mkString
            newBuf append code
          } else newBuf += next
          newBuf append LiteralColor
          closing = 0
        } else if (shouldInterpolate) {
          interpolate()
        } else if (curr == delim && multiline) {
          closing += 1
          if (closing == 3) continue = false
          newBuf += curr
        } else if (curr == delim) {
          continue = false
          newBuf += curr
        } else {
          newBuf += curr
          closing = 0
        }

        if (Chars.isLineBreakChar(curr)) {
          newBuf append LiteralColor
        }
      }
      newBuf append NoColor
      prev = curr
    }

    def appendSingleQuote(delim: Char) = remaining.take(3) match {
      case chr #:: '\'' #:: _ => // single character
        newBuf append LiteralColor
        newBuf appendAll s"'$chr'"
        newBuf append NoColor
        takeChars(2)
        prev = '\''
      case '\\' #:: chr #:: '\'' #:: _ => // escaped character
        newBuf append LiteralColor
        newBuf appendAll s"'\\$chr'"
        newBuf append NoColor
        takeChars(3)
        prev = '\''
      case _ => appendWhile(delim, !typeEnders.contains(_), literal)
    }

    def append(c: Char, shouldHL: String => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")

      def delim(c: Char) = (c: @switch) match {
        case ' ' => true
        case '\n' => true
        case '(' => true
        case ')' => true
        case '[' => true
        case ']' => true
        case ':' => true
        case '@' => true
        case ',' => true
        case '.' => true
        case _ => false
      }

      val valDefStarterTokens = "var" :: "val" :: "def" :: "case" :: Nil

      /** lastValDefToken is used to check whether we want to show something
       *  in valDef color or not. There are only a few cases when lastValDefToken
       *  should be updated, that way we can avoid stopping coloring too early.
       *  eg.: case A(x, y, z) => ???
       *  Without this function only x would be colored.
       */
      def updateLastToken(currentToken: String): String =
        (lastValDefToken, currentToken) match {
          case _ if valDefStarterTokens.contains(currentToken) => currentToken
          case (("val" | "var"), "=") => currentToken
          case ("case", ("=>" | "class" | "object")) => currentToken
          case ("def", _) => currentToken
          case _ => lastValDefToken
        }

      while (remaining.nonEmpty && !delim(curr)) {
        curr = takeChar()
        if (!delim(curr)) sb += curr
      }

      val str    = sb.toString
      val toAdd  =
        if (shouldHL(str))
          highlight(str)
        else if (valDefStarterTokens.contains(lastValDefToken) && !List("=", "=>").contains(str))
          valDef(str)
        else str
      val suffix = if (delim(curr)) s"$curr" else ""
      newBuf append (toAdd + suffix)
      lastValDefToken = updateLastToken(str)
      prev = curr
    }

    def appendWhile(c: Char, pred: Char => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")
      while (remaining.nonEmpty && pred(curr)) {
        curr = takeChar()
        if (pred(curr)) sb += curr
      }

      val str    = sb.toString
      val suffix = if (!pred(curr)) s"$curr" else ""
      newBuf append (highlight(str) + suffix)
      prev = curr
    }

    def appendTo(c: Char, pred: Char => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")
      while (remaining.nonEmpty && !pred(curr)) {
        curr = takeChar()
        sb += curr
      }

      newBuf append highlight(sb.toString)
      prev = curr
    }

    newBuf.toIterable
  }

  private class NoReporter extends Reporter {
    override def doReport(m: MessageContainer)(implicit ctx: Context): Unit = ()
  }

  def highlight(in: String)(ctx0: Context): String = {
    import dotty.tools.dotc.ast.untpd._

    implicit val ctx: Context = ctx0.fresh.setReporter(new NoReporter)

    val sf = new SourceFile("<highlighting>", in.toCharArray)
    val p = new Parser(sf)
    val s = new Scanner(sf)
    val trees = p.blockStatSeq()

    val outputH = Array.fill(in.length)(NoColor)

    def highlightRange(p: Position, color: String): Unit = {
      if(p.exists) {
        for {
          i <- p.start until math.min(p.end, outputH.length)
        } outputH(i) = color
      }
    }

    val treeTraverser = new UntypedTreeTraverser {
      def traverse(tree: Tree)(implicit ctx: Context): Unit = {
        tree match {
          case tpe : TypeDef =>
            highlightRange(tpe.namePos, TypeColor)
          case _ : TypTree =>
            highlightRange(tree.pos, TypeColor)
          case mod: ModuleDef =>
            highlightRange(mod.namePos, TypeColor)
          case v : ValOrDefDef =>
            highlightRange(v.namePos, ValDefColor)
            highlightRange(v.tpt.pos, TypeColor)
          case _ : Literal =>
            highlightRange(tree.pos, LiteralColor)
          case _ =>
        }
        traverseChildren(tree)
      }
    }

    for {
      t <- trees
    } {
      treeTraverser.traverse(t)
    }

    val sb = new mutable.StringBuilder()

    while(s.token != EOF) {
      val isKwd = isKeyword(s.token)
      val offsetStart = s.offset


      if(s.token == IDENTIFIER && s.name == nme.???) {
        highlightRange(Position(s.offset, s.offset + s.name.length), Console.RED_B)
      }
      s.nextToken()

      if(isKwd) {
        val offsetEnd = s.lastOffset
        highlightRange(Position(offsetStart, offsetEnd), KeywordColor)
      }
    }

    for {
      idx <- outputH.indices
    } {
      if(idx == 0 || outputH(idx-1) != outputH(idx)){
        sb.append(outputH(idx))
      }
      sb.append(in(idx))
    }
    sb.append(NoColor)

    sb.mkString
  }
}
<commit_msg>Small improvements in highlighter<commit_after>package dotty.tools
package dotc
package printing

import dotty.tools.dotc.core.Contexts.Context
import dotty.tools.dotc.core.StdNames._
import dotty.tools.dotc.parsing.Parsers.Parser
import dotty.tools.dotc.parsing.Scanners.Scanner
import dotty.tools.dotc.parsing.Tokens._
import dotty.tools.dotc.reporting.Reporter
import dotty.tools.dotc.reporting.diagnostic.MessageContainer
import dotty.tools.dotc.util.Positions.Position

import scala.annotation.switch
import scala.collection.mutable.StringBuilder
import util.{Chars, SourceFile}

import scala.collection.mutable

/** This object provides functions for syntax highlighting in the REPL */
object SyntaxHighlighting {

  // Keep in sync with SyntaxHighlightingTests
  val NoColor         = Console.RESET
  val CommentColor    = Console.BLUE
  val KeywordColor    = Console.YELLOW
  val ValDefColor     = Console.CYAN
  val LiteralColor    = Console.RED
  val StringColor     = Console.GREEN
  val TypeColor       = Console.MAGENTA
  val AnnotationColor = Console.MAGENTA

  private def none(str: String) = str
  private def keyword(str: String) = KeywordColor + str + NoColor
  private def typeDef(str: String) = TypeColor + str + NoColor
  private def literal(str: String) = LiteralColor + str + NoColor
  private def valDef(str: String) = ValDefColor + str + NoColor
  private def operator(str: String) = TypeColor + str + NoColor
  private def annotation(str: String) =
    if (str.trim == "@") str else { AnnotationColor + str + NoColor }
  private val tripleQs = Console.RED_B + "???" + NoColor

  private val keywords: Seq[String] = for {
    index <- IF to ERASED // All alpha keywords
  } yield tokenString(index)

  private val interpolationPrefixes =
    'A' :: 'B' :: 'C' :: 'D' :: 'E' :: 'F' :: 'G' :: 'H' :: 'I' :: 'J' :: 'K' ::
    'L' :: 'M' :: 'N' :: 'O' :: 'P' :: 'Q' :: 'R' :: 'S' :: 'T' :: 'U' :: 'V' ::
    'W' :: 'X' :: 'Y' :: 'Z' :: '$' :: '_' :: 'a' :: 'b' :: 'c' :: 'd' :: 'e' ::
    'f' :: 'g' :: 'h' :: 'i' :: 'j' :: 'k' :: 'l' :: 'm' :: 'n' :: 'o' :: 'p' ::
    'q' :: 'r' :: 's' :: 't' :: 'u' :: 'v' :: 'w' :: 'x' :: 'y' :: 'z' :: Nil

  private val typeEnders =
   '{'  :: '}' :: ')' :: '(' :: '[' :: ']' :: '=' :: ' ' :: ',' :: '.' :: '|' ::
   '&' :: '\n' :: Nil

  def apply(chars: Iterable[Char])(implicit ctx: Context): Iterable[Char] = {
    if (ctx.settings.color.value != "never") highlight(chars)
    else chars
  }

  def highlight(chars: Iterable[Char]): Iterable[Char] = {
    var prev: Char = 0
    var remaining  = chars.toStream
    val newBuf     = new StringBuilder
    var lastValDefToken  = ""

    @forceInline def keywordStart =
      prev == 0    || prev == ' ' || prev == '{' || prev == '(' ||
      prev == '\n' || prev == '[' || prev == ',' || prev == ':' ||
      prev == '|'  || prev == '&' || prev.isDigit

    @forceInline def numberStart(c: Char) =
      c.isDigit && (!prev.isLetter || prev == '.' || prev == ' ' || prev == '(' || prev == '\u0000')

    def takeChar(): Char = takeChars(1).head
    def takeChars(x: Int): Seq[Char] = {
      val taken = remaining.take(x)
      remaining = remaining.drop(x)
      taken
    }

    while (remaining.nonEmpty) {
      val n = takeChar()
      if (interpolationPrefixes.contains(n)) {
        // Interpolation prefixes are a superset of the keyword start chars
        val (prefix, after) = remaining.span(interpolationPrefixes.contains)
        if (after.startsWith("\"")) {
          newBuf += n ++= prefix
          prev = prefix.lastOption.getOrElse(n)
          if (remaining.nonEmpty) takeChars(prefix.length + 1) // drop 1 for appendLiteral
          appendString('"', after.startsWith("\"\"\""), true)
        } else {
          if (n.isUpper && (keywordStart || prev == '.')) {
            appendWhile(n, !typeEnders.contains(_), typeDef)
          } else if (keywordStart) {
            append(n, keywords.contains(_), { kw =>
              if (kw == "new") typeDef(kw) else keyword(kw)
            })
          } else {
            newBuf += n
            prev = n
          }
        }
      } else {
        (n: @switch) match {
          case '/'  =>
            if (remaining.nonEmpty) {
              remaining.head match {
                case '/' =>
                  takeChar()
                  eolComment()
                case '*' =>
                  takeChar()
                  blockComment()
                case x =>
                  newBuf += '/'
              }
            } else newBuf += '/'
          case '='  =>
            append('=', _ == "=>", operator)
          case '<'  =>
            append('<', { x => x == "<-" || x == "<:" || x == "<%" }, operator)
          case '>'  =>
            append('>', { x => x == ">:" }, operator)
          case '#'  =>
            if (prev != ' ' && prev != '.') newBuf append operator("#")
            else newBuf += n
            prev = '#'
          case '@'  =>
            appendWhile('@', !typeEnders.contains(_), annotation)
          case '\"' =>
            appendString('\"', multiline = remaining.take(2).mkString == "\"\"", false)
          case '\'' =>
            appendSingleQuote('\'')
          case '`'  =>
            appendTo('`', _ == '`', none)
          case _    => {
            if (n == '?' && remaining.take(2).mkString == "??") {
              takeChars(2)
              newBuf append tripleQs
              prev = '?'
            }
            else if (n.isUpper && keywordStart)
              appendWhile(n, !typeEnders.contains(_), typeDef)
            else if (numberStart(n)) {
              def isNumber(c: Char): Boolean =
                c.isDigit || c == '\u0000' || (c == '.' && remaining.nonEmpty && remaining.head.isDigit)
              appendWhile(n, isNumber, literal)
            } else
              newBuf += n; prev = n
          }
        }
      }
    }

    def eolComment() = {
      newBuf append (CommentColor + "//")
      var curr = '/'
      while (curr != '\n' && remaining.nonEmpty) {
        curr = takeChar()
        newBuf += curr
      }
      prev = curr
      newBuf append NoColor
    }

    def blockComment() = {
      newBuf append (CommentColor + "/*")
      var curr = '*'
      var open = 1
      while (open > 0 && remaining.nonEmpty) {
        curr = takeChar()
        if (curr == '@') {
          appendWhile('@', !typeEnders.contains(_), annotation)
          newBuf append CommentColor
        }
        else newBuf += curr

        if (curr == '*' && remaining.nonEmpty) {
          curr = takeChar()
          newBuf += curr
          if (curr == '/') open -= 1
        } else if (curr == '/' && remaining.nonEmpty) {
          curr = takeChar()
          newBuf += curr
          if (curr == '*') open += 1
        }

        if (Chars.isLineBreakChar(curr)) {
          newBuf append CommentColor
        }
      }
      prev = curr
      newBuf append NoColor
    }

    def appendString(delim: Char, multiline: Boolean = false, inInterpolation: Boolean) = {
      var curr: Char      = 0
      var continue        = true
      var closing         = 0
      newBuf append (LiteralColor + delim)

      def shouldInterpolate =
        inInterpolation && curr == '$' && prev != '$' && remaining.nonEmpty

      def interpolate() = {
        val next = takeChar()
        if (next == '$') {
          newBuf += curr
          newBuf += next
          prev = '$'
        } else if (next == '{') {
          var open = 1 // keep track of open blocks
          newBuf append (ValDefColor + curr)
          newBuf += next
          while (remaining.nonEmpty && open > 0) {
            var c = takeChar()
            newBuf += c
            if (c == '}') open -= 1
            else if (c == '{') open += 1
          }
          newBuf append LiteralColor
        } else {
          newBuf append (ValDefColor + curr)
          newBuf += next
          var c: Char = 'a'
          while (c.isLetterOrDigit && remaining.nonEmpty) {
            c = takeChar()
            if (c != '"') newBuf += c
          }
          newBuf append LiteralColor
          if (c == '"') {
            newBuf += c
            continue = false
          }
        }
        closing = 0
      }

      while (continue && remaining.nonEmpty) {
        curr = takeChar()
        if (curr == '\\' && remaining.nonEmpty) {
          val next = takeChar()
          newBuf append (KeywordColor + curr)
          if (next == 'u') {
            val code = "u" + takeChars(4).mkString
            newBuf append code
          } else newBuf += next
          newBuf append LiteralColor
          closing = 0
        } else if (shouldInterpolate) {
          interpolate()
        } else if (curr == delim && multiline) {
          closing += 1
          if (closing == 3) continue = false
          newBuf += curr
        } else if (curr == delim) {
          continue = false
          newBuf += curr
        } else {
          newBuf += curr
          closing = 0
        }

        if (Chars.isLineBreakChar(curr)) {
          newBuf append LiteralColor
        }
      }
      newBuf append NoColor
      prev = curr
    }

    def appendSingleQuote(delim: Char) = remaining.take(3) match {
      case chr #:: '\'' #:: _ => // single character
        newBuf append LiteralColor
        newBuf appendAll s"'$chr'"
        newBuf append NoColor
        takeChars(2)
        prev = '\''
      case '\\' #:: chr #:: '\'' #:: _ => // escaped character
        newBuf append LiteralColor
        newBuf appendAll s"'\\$chr'"
        newBuf append NoColor
        takeChars(3)
        prev = '\''
      case _ => appendWhile(delim, !typeEnders.contains(_), literal)
    }

    def append(c: Char, shouldHL: String => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")

      def delim(c: Char) = (c: @switch) match {
        case ' ' => true
        case '\n' => true
        case '(' => true
        case ')' => true
        case '[' => true
        case ']' => true
        case ':' => true
        case '@' => true
        case ',' => true
        case '.' => true
        case _ => false
      }

      val valDefStarterTokens = "var" :: "val" :: "def" :: "case" :: Nil

      /** lastValDefToken is used to check whether we want to show something
       *  in valDef color or not. There are only a few cases when lastValDefToken
       *  should be updated, that way we can avoid stopping coloring too early.
       *  eg.: case A(x, y, z) => ???
       *  Without this function only x would be colored.
       */
      def updateLastToken(currentToken: String): String =
        (lastValDefToken, currentToken) match {
          case _ if valDefStarterTokens.contains(currentToken) => currentToken
          case (("val" | "var"), "=") => currentToken
          case ("case", ("=>" | "class" | "object")) => currentToken
          case ("def", _) => currentToken
          case _ => lastValDefToken
        }

      while (remaining.nonEmpty && !delim(curr)) {
        curr = takeChar()
        if (!delim(curr)) sb += curr
      }

      val str    = sb.toString
      val toAdd  =
        if (shouldHL(str))
          highlight(str)
        else if (valDefStarterTokens.contains(lastValDefToken) && !List("=", "=>").contains(str))
          valDef(str)
        else str
      val suffix = if (delim(curr)) s"$curr" else ""
      newBuf append (toAdd + suffix)
      lastValDefToken = updateLastToken(str)
      prev = curr
    }

    def appendWhile(c: Char, pred: Char => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")
      while (remaining.nonEmpty && pred(curr)) {
        curr = takeChar()
        if (pred(curr)) sb += curr
      }

      val str    = sb.toString
      val suffix = if (!pred(curr)) s"$curr" else ""
      newBuf append (highlight(str) + suffix)
      prev = curr
    }

    def appendTo(c: Char, pred: Char => Boolean, highlight: String => String) = {
      var curr: Char = 0
      val sb = new StringBuilder(s"$c")
      while (remaining.nonEmpty && !pred(curr)) {
        curr = takeChar()
        sb += curr
      }

      newBuf append highlight(sb.toString)
      prev = curr
    }

    newBuf.toIterable
  }

  private class NoReporter extends Reporter {
    override def doReport(m: MessageContainer)(implicit ctx: Context): Unit = ()
  }

  private val ignoredKwds = Seq(nme.ARROWkw, nme.EQ, nme.EQL, nme.COLONkw)

  def highlight(in: String)(ctx0: Context): String = {
    import dotty.tools.dotc.ast.untpd._

    implicit val ctx: Context = ctx0.fresh.setReporter(new NoReporter)

    val sf = new SourceFile("<highlighting>", in.toCharArray)
    val p = new Parser(sf)
    val s = new Scanner(sf)
    val trees = p.blockStatSeq()

    val outputH = Array.fill(in.length)(NoColor)

    def highlightRange(p: Position, color: String): Unit = {
      if(p.exists) {
        for {
          i <- p.start until math.min(p.end, outputH.length)
        } outputH(i) = color
      }
    }

    val treeTraverser = new UntypedTreeTraverser {
      def traverse(tree: Tree)(implicit ctx: Context): Unit = {
        tree match {
          case tpe : TypeDef =>
            highlightRange(tpe.namePos, TypeColor)
          case _ : TypTree =>
            highlightRange(tree.pos, TypeColor)
          case mod: ModuleDef =>
            highlightRange(mod.namePos, TypeColor)
          case v : ValOrDefDef =>
            highlightRange(v.namePos, ValDefColor)
            highlightRange(v.tpt.pos, TypeColor)
          case _ : Literal =>
            highlightRange(tree.pos, LiteralColor)
          case _ =>
        }
        traverseChildren(tree)
      }
    }

    for {
      t <- trees
    } {
      treeTraverser.traverse(t)
    }

    val sb = new mutable.StringBuilder()

    while(s.token != EOF) {
      val isKwd = isKeyword(s.token) && !ignoredKwds.contains(s.name)
      val offsetStart = s.offset

      if(s.token == IDENTIFIER && s.name == nme.???) {
        highlightRange(Position(s.offset, s.offset + s.name.length), Console.RED_B)
      }
      s.nextToken()

      if(isKwd) {
        val offsetEnd = s.lastOffset
        highlightRange(Position(offsetStart, offsetEnd), KeywordColor)
      }
    }

    for {
      idx <- outputH.indices
    } {
      if(idx == 0 || outputH(idx-1) != outputH(idx)){
        sb.append(outputH(idx))
      }
      sb.append(in(idx))
    }
    if(outputH.last != NoColor) {
      sb.append(NoColor)
    }

    sb.mkString
  }
}

{'fraction_non_alphanumeric': 0.14780265339966833, 'fraction_numerical': 0.002902155887230514, 'mean_word_length': 1.8344104974539757, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 117, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4554460', 'n_tokens_mistral': 8999, 'n_tokens_neox': 8798, 'n_words': 2241}