Ensure root is set in both worlds. Fixes #159.

Author: wmdietl

src/main/java/com/google/jspecify/nullness/NullSpecAnnotatedTypeFactory.java

@@ -43,6 +43,7 @@
 
 import com.sun.source.tree.BinaryTree;
 import com.sun.source.tree.ClassTree;
+import com.sun.source.tree.CompilationUnitTree;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.IdentifierTree;
 import com.sun.source.tree.LambdaExpressionTree;
@@ -374,6 +375,27 @@ private NullSpecAnnotatedTypeFactory(
     }
   }
 
+  /** Too ensure setRoot is called on both worlds exactly once. */
+  private boolean settingRoot = false;
+
+  /**
+   * Ensure setRoot is called on both worlds exactly once whenever it is called on one of the
+   * worlds.
+   */
+  @Override
+  public void setRoot(@Nullable CompilationUnitTree root) {
+    if (!settingRoot) {
+      settingRoot = true;
+      super.setRoot(root);
+      if (withLeastConvenientWorld != this) {
+        withLeastConvenientWorld.setRoot(root);
+      } else {
+        withMostConvenientWorld.setRoot(root);
+      }
+    }
+    settingRoot = false;
+  }
+
   @Override
   protected void addCheckedCodeDefaults(QualifierDefaults defs) {
     // TODO: add false for subpackages once overload is added to CF. Shouldn't really matter.