Rewrite imports collection, avoid deprecated code

Author: kitbellew

scalafix-core/src/main/scala/scalafix/internal/patch/ImportPatchOps.scala

@@ -40,10 +40,11 @@ object ImportPatchOps {
     }
   }
   private def fallbackToken(ctx: RuleCtx): Token = {
+    @tailrec
     def loop(tree: Tree): Token = tree match {
       case Source((stat: Pkg) :: _) => loop(stat)
-      case Source(_) => tree.tokens.head
-      case Pkg(_, stat :: _) => loop(stat)
+      case _: Source => tree.tokens.head
+      case Pkg.Initial(_, stat :: _) => loop(stat)
       case els =>
         ctx.tokenList.prev(ctx.tokenList.prev(els.tokens.head)) match {
           case comment @ Token.Comment(_) =>
@@ -53,19 +54,12 @@ object ImportPatchOps {
     }
     loop(ctx.tree)
   }
-  private def extractImports(stats: Seq[Stat]): Seq[Import] = {
-    stats
-      .takeWhile(_.is[Import])
-      .collect { case i: Import => i }
-  }
 
-  @tailrec private final def getGlobalImports(ast: Tree): Seq[Import] =
-    ast match {
-      case Pkg(_, Seq(pkg: Pkg)) => getGlobalImports(pkg)
-      case Source(Seq(pkg: Pkg)) => getGlobalImports(pkg)
-      case Pkg(_, stats) => extractImports(stats)
-      case Source(stats) => extractImports(stats)
-      case _ => Nil
+  @tailrec
+  private def extractImports(stats: List[Stat]): Seq[Import] =
+    stats match {
+      case (pkg: Pkg) :: Nil => extractImports(pkg.body.stats)
+      case _ => stats.takeWhile(_.is[Import]).collect { case i: Import => i }
     }
 
   // NOTE(olafur): This method is the simplest/dummest thing I can think of
@@ -88,7 +82,11 @@ object ImportPatchOps {
     lazy val allImporteeSymbols = allImportees.flatMap(importee =>
       importee.symbol.map(_.normalized -> importee)
     )
-    val globalImports = getGlobalImports(ctx.tree)
+    val globalImports = ctx.tree match {
+      case t: Source => extractImports(t.stats)
+      case t: Pkg => extractImports(t.body.stats)
+      case _ => Nil
+    }
     val editToken: Token = {
       if (globalImports.isEmpty) fallbackToken(ctx)
       else globalImports.last.tokens.last

scalafix-core/src/main/scala/scalafix/internal/patch/ReplaceSymbolOps.scala

@@ -22,25 +22,22 @@ object ReplaceSymbolOps {
     }
   }
 
-  private def extractImports(stats: Seq[Stat]): Seq[Import] = {
-    stats.collect { case i: Import => i }
+  @tailrec
+  private def extractImports(stats: List[Stat]): Seq[Import] = stats match {
+    case (p: Pkg) :: Nil => extractImports(p.body.stats)
+    case _ => stats.collect { case i: Import => i }
   }
 
   private def getNamesOfExplicitlyImportedSymbols(
       tree: Tree,
       isMoved: Name => Boolean
   ): Set[String] = {
-    @tailrec
-    def getGlobalImports(ast: Tree): Seq[Import] = ast match {
-      case Pkg(_, Seq(pkg: Pkg)) => getGlobalImports(pkg)
-      case Source(Seq(pkg: Pkg)) => getGlobalImports(pkg)
-      case Pkg(_, stats) => extractImports(stats)
-      case Source(stats) => extractImports(stats)
+    val globalImports = tree match {
+      case t: Pkg => extractImports(t.body.stats)
+      case t: Source => extractImports(t.stats)
       case _ => Nil
     }
 
-    val globalImports = getGlobalImports(tree)
-
     // pre-compute global imported symbols for O(1) collision detection
     // since ctx.addGlobalImport adds imports at global scope
     // exclude names whose symbols are being moved, as those imports

scalafix-rules/src/main/scala/scalafix/internal/rule/OrganizeImports.scala

@@ -76,7 +76,11 @@ class OrganizeImports(
 
     val unusedImporteePositions = new UnusedImporteePositions
 
-    val (globalImports, localImports) = collectImports(doc.tree)
+    val (globalImports, localImports) = doc.tree match {
+      case t: Source => collectImports(t.stats)
+      case t: Pkg => collectImports(t.body.stats)
+      case _ => (Nil, Nil)
+    }
 
     val globalImportsPatch =
       if (globalImports.isEmpty) Patch.empty
@@ -878,22 +882,22 @@ object OrganizeImports {
     }
 
   @tailrec private def collectImports(
-      tree: Tree
-  ): (Seq[Import], Seq[Import]) = {
-    def extractImports(stats: Seq[Stat]): (Seq[Import], Seq[Import]) = {
-      val (importStats, otherStats) = stats.span(_.is[Import])
-      val globalImports = importStats.map { case i: Import => i }
-      val localImports = otherStats.flatMap(_.collect { case i: Import => i })
-      (globalImports, localImports)
-    }
-
-    tree match {
-      case Source(Seq(p: Pkg)) => collectImports(p)
-      case Pkg(_, Seq(p: Pkg)) => collectImports(p)
-      case Source(stats) => extractImports(stats)
-      case Pkg(_, stats) => extractImports(stats)
-      case _ => (Nil, Nil)
-    }
+      stats: List[Stat]
+  ): (Seq[Import], Seq[Import]) = stats match {
+    case (p: Pkg) :: Nil => collectImports(p.body.stats)
+    case _ =>
+      val globalImports = Seq.newBuilder[Import]
+      val localImports = Seq.newBuilder[Import]
+      def collectLocalImports(tree: Tree): Unit =
+        tree.traverse { case i: Import => localImports += i }
+      val statsiter = stats.iterator
+      while (statsiter.hasNext) statsiter.next() match {
+        case i: Import => globalImports += i
+        case i =>
+          collectLocalImports(i)
+          while (statsiter.hasNext) collectLocalImports(statsiter.next())
+      }
+      (globalImports.result(), localImports.result())
   }
 
   @tailrec private def topQualifierOf(term: Term): Term.Name =