Fix #412: Protect the state of packages and loaders from exceptions.

tasty-query/shared/src/main/scala/tastyquery/Symbols.scala

@@ -1801,6 +1801,8 @@ object Symbols {
     // DeclaringSymbol-related fields
     private var rootsInitialized: Boolean = false
     private val myDeclarations = mutable.HashMap[UnsignedName, Symbol]()
+    private val pendingDeclarations = mutable.HashMap[UnsignedName, Symbol]()
+    private var isLoadingNewRoots: Boolean = false
 
     // Cache fields
     val packageRef: PackageRef = new PackageRef(this)
@@ -1852,19 +1854,58 @@ object Symbols {
     end getPackageDecl
 
     private[Symbols] final def addDecl(decl: Symbol): Unit =
-      assert(!myDeclarations.contains(decl.name), s"trying to add a second entry $decl for name ${decl.name} in $this")
-      myDeclarations(decl.name) = decl
+      assert(
+        !myDeclarations.contains(decl.name) && !pendingDeclarations.contains(decl.name),
+        s"trying to add a second entry $decl for name ${decl.name} in $this"
+      )
+
+      /* If we are loading new roots and the decl is not a package,
+       * add the declaration to the pending set only. They will be committed
+       * later by `loadingNewRoots` if loading is successful.
+       *
+       * Packages are always eagerly committed.
+       */
+      decl match
+        case decl: TermOrTypeSymbol if isLoadingNewRoots =>
+          pendingDeclarations(decl.name) = decl
+        case _ =>
+          myDeclarations(decl.name) = decl
+    end addDecl
+
+    /** Performs an operation that can load new roots from the class loader.
+      *
+      * While loading new roots, any new non-package member sent to `addDecl`
+      * is added to `pendingDeclarations` instead of `myDeclarations`. They
+      * are committed to `myDeclarations` only after the `op` successfully
+      * completes.
+      *
+      * This way, any exception occurring during loading does not pollute the
+      * publicly visible state in `myDeclarations`.
+      */
+    private def loadingNewRoots[A](op: Loader => A)(using Context): A =
+      if isLoadingNewRoots then throw IllegalStateException(s"Cyclic loading of new roots in package $this")
+
+      isLoadingNewRoots = true
+      try
+        if !rootsInitialized then
+          ctx.classloader.scanPackage(this)
+          rootsInitialized = true
+
+        val result = op(ctx.classloader)
+
+        // Upon success, commit pending declations
+        myDeclarations ++= pendingDeclarations
 
-    private final def ensureRootsInitialized()(using Context): Unit =
-      if !rootsInitialized then
-        ctx.classloader.scanPackage(this)
-        rootsInitialized = true
+        result
+      finally
+        pendingDeclarations.clear() // whether or not they were committed
+        isLoadingNewRoots = false
+    end loadingNewRoots
 
     final def getDecl(name: Name)(using Context): Option[Symbol] = name match
       case name: UnsignedName =>
         myDeclarations.get(name).orElse {
-          ensureRootsInitialized()
-          if ctx.classloader.loadRoot(this, name) then myDeclarations.get(name)
+          if loadingNewRoots(_.loadRoot(this, name)) then myDeclarations.get(name)
           else None
         }
       case _: SignedName =>
@@ -1893,8 +1934,7 @@ object Symbols {
       }
 
     final def declarations(using Context): List[Symbol] =
-      ensureRootsInitialized()
-      ctx.classloader.loadAllRoots(this)
+      loadingNewRoots(_.loadAllRoots(this))
       myDeclarations.values.toList
 
     // See PackageRef.findMember
@@ -1908,8 +1948,7 @@ object Symbols {
     end allPackageObjectDecls
 
     private def computeAllPackageObjectDecls()(using Context): List[ClassSymbol] =
-      ensureRootsInitialized()
-      ctx.classloader.loadAllPackageObjectRoots(this)
+      loadingNewRoots(_.loadAllPackageObjectRoots(this))
       myDeclarations.valuesIterator.collect {
         case cls: ClassSymbol if cls.name.isPackageObjectClassName => cls
       }.toList

tasty-query/shared/src/main/scala/tastyquery/reader/Loaders.scala

@@ -159,10 +159,15 @@ private[tastyquery] object Loaders {
 
     /** Loads all the roots of the given `pkg`. */
     private[tastyquery] def loadAllRoots(pkg: PackageSymbol)(using Context): Unit =
-      for
-        entries <- roots.remove(pkg)
-        (rootName, entry) <- IArray.from(entries).sortBy(_(0).name) // sort for determinism.
-      do completeRoot(Loader.Root(pkg, rootName), entry)
+      roots.get(pkg) match
+        case Some(entries) =>
+          val allNames =
+            entries.keysIterator.toList.sortBy(_.name) // sort for determinism
+          for rootName <- allNames do doLoadRoot(pkg, entries, rootName)
+          // Upon success, we won't need anything from that package anymore
+          roots -= pkg
+        case None =>
+          ()
     end loadAllRoots
 
     /** Loads all the roots of the given `pkg` that could be package objects. */
@@ -171,10 +176,9 @@ private[tastyquery] object Loaders {
         case Some(entries) =>
           val candidateNames =
             entries.keysIterator.filter(_.isPackageObjectName).toList.sortBy(_.name) // sort for determinism
-          for
-            rootName <- candidateNames
-            entry <- entries.remove(rootName)
-          do completeRoot(Loader.Root(pkg, rootName), entry)
+          for rootName <- candidateNames do doLoadRoot(pkg, entries, rootName)
+          // Upon success, we won't need any of the loaded entries anymore
+          entries --= candidateNames
         case None =>
           ()
     end loadAllPackageObjectRoots
@@ -196,16 +200,26 @@ private[tastyquery] object Loaders {
       roots.get(pkg) match
         case Some(entries) =>
           val rootName = topLevelSymbolNameToRootName(name)
-          entries.remove(rootName) match
-            case Some(entry) =>
-              completeRoot(Loader.Root(pkg, rootName), entry)
-              true
-            case None =>
-              false
+          val result = doLoadRoot(pkg, entries, rootName)
+          if result then
+            // Upon success, we won't need that particular entry anymore
+            entries -= rootName
+          result
         case None =>
           false
     end loadRoot
 
+    private def doLoadRoot(pkg: PackageSymbol, pkgEntries: mutable.Map[SimpleName, Entry], rootName: SimpleName)(
+      using Context
+    ): Boolean =
+      pkgEntries.get(rootName) match
+        case Some(entry) =>
+          completeRoot(Loader.Root(pkg, rootName), entry)
+          true
+        case None =>
+          false
+    end doLoadRoot
+
     private def foreachEntry(data: PackageData)(f: (SimpleName, Entry) => Unit): Unit =
       def binaryNameToRootName(binaryName: String): SimpleName =
         termName(NameTransformer.decode(binaryName))
@@ -239,11 +253,11 @@ private[tastyquery] object Loaders {
       require(searched)
       packages.get(pkg) match {
         case Some(datas) =>
-          packages -= pkg
           val localRoots = mutable.HashMap.empty[SimpleName, Entry]
           for data <- datas do
             foreachEntry(data)(localRoots.getOrElseUpdate) // duplicate roots from later classpath entries are ignored
           roots(pkg) = localRoots
+          packages -= pkg
         case _ => // probably a synthetic package that only has other packages as members. (i.e. `package java`)
       }
     }