gominimal · dependabot · Sep 12, 2025 · Sep 26, 2025
diff --git a/cli/src/doctest.rs b/cli/src/doctest.rs
@@ -296,7 +296,7 @@ impl TestCommand {
         let mut registry = TestRegistry::default();
         program.typecheck(TypecheckMode::Walk)?;
         program
-            .custom_transform(|cache, rt| doctest_transform(cache, &mut registry, rt))
+            .custom_transform(0, |cache, rt| doctest_transform(cache, &mut registry, rt))
             .map_err(|e| e.unwrap_error("transforming doctest"))?;
         Ok((program.eval_closurized_record_spine()?, registry))
     }

diff --git a/core/src/cache.rs b/core/src/cache.rs
@@ -1037,40 +1037,46 @@ impl CacheHub {
         Ok(())
     }
 
-    /// Applies a custom transform to an input and its imports, leaving them in the same state as
-    /// before. Requires that the input has been parsed. In order for the transform to apply to
-    /// imports, they need to have been resolved.
+    /// Applies a custom transform to an input and its imports. [CacheError::NotTypechecked] is returned
+    /// if the file has not yet been typechecked.
+    ///
+    /// If multiple invocations of `custom_transform` are needed, you must supply `transform_id` with
+    /// with a number higher than that of all previous invocations.
     pub fn custom_transform<E>(
         &mut self,
         file_id: FileId,
+        transform_id: usize,
         f: &mut impl FnMut(&mut CacheHub, RichTerm) -> Result<RichTerm, E>,
     ) -> Result<(), CacheError<E>> {
         match self.terms.entry_state(file_id) {
-            Some(state) if state >= EntryState::Parsed => {
-                if state < EntryState::Transforming {
+            None => Err(CacheError::NotParsed),
+            Some(state) if state < EntryState::Typechecked => Err(CacheError::NotTypechecked),
+            Some(state) => {
+                if state.needs_custom_transform(transform_id) {
                     let cached_term = self.terms.terms.remove(&file_id).unwrap();
                     let term = f(self, cached_term.term)?;
                     self.terms.insert(
                         file_id,
                         TermEntry {
                             term,
-                            state: EntryState::Transforming,
+                            state: EntryState::CustomTransforming,
                             ..cached_term
                         },
                     );
 
                     if let Some(imports) = self.import_data.imports.get(&file_id).cloned() {
                         for file_id in imports.into_iter() {
-                            self.custom_transform(file_id, f)?;
+                            self.custom_transform(file_id, transform_id, f)?;
                         }
                     }
-                    // TODO: We're setting the state back to whatever it was.
                     // unwrap(): we inserted the term just above
-                    let _ = self.terms.update_state(file_id, state).unwrap();
+                    let _ = self
+                        .terms
+                        .update_state(file_id, EntryState::CustomTransformed { transform_id })
+                        .unwrap();
                 }
                 Ok(())
             }
-            _ => Err(CacheError::NotParsed),
         }
     }
 
@@ -1447,6 +1453,10 @@ pub enum EntryState {
     Typechecking,
     /// The entry and its transitive imports have been typechecked.
     Typechecked,
+    /// A custom transformation of the entry (through `Program::custom_transform`) is underway.
+    CustomTransforming,
+    /// This entry has completed custom transformations of this ID and lower.
+    CustomTransformed { transform_id: usize },
     /// The imports of the entry have been resolved, and the imports of its (transitive) imports are
     /// being resolved.
     ImportsResolving,
@@ -1460,6 +1470,19 @@ pub enum EntryState {
     Closurized,
 }
 
+impl EntryState {
+    fn needs_custom_transform(&self, transform_id: usize) -> bool {
+        if let EntryState::CustomTransformed {
+            transform_id: done_transform_id,
+        } = self
+        {
+            transform_id > *done_transform_id
+        } else {
+            self < &EntryState::CustomTransforming
+        }
+    }
+}
+
 /// The result of a cache operation, such as parsing, typechecking, etc. which can either have
 /// performed actual work, or have done nothing if the corresponding entry was already at a later
 /// stage.
@@ -1483,6 +1506,7 @@ impl<T> CacheOp<T> {
 pub enum CacheError<E> {
     Error(E),
     NotParsed,
+    NotTypechecked,
 }
 
 impl<E> From<E> for CacheError<E> {
@@ -1496,14 +1520,15 @@ impl<E> CacheError<E> {
     pub fn unwrap_error(self, msg: &str) -> E {
         match self {
             CacheError::Error(err) => err,
-            CacheError::NotParsed => panic!("{}", msg),
+            CacheError::NotParsed | CacheError::NotTypechecked => panic!("{}", msg),
         }
     }
 
     pub fn map_err<O>(self, f: impl FnOnce(E) -> O) -> CacheError<O> {
         match self {
             CacheError::Error(e) => CacheError::Error(f(e)),
             CacheError::NotParsed => CacheError::NotParsed,
+            CacheError::NotTypechecked => CacheError::NotTypechecked,
         }
     }
 }

diff --git a/core/src/program.rs b/core/src/program.rs
@@ -541,18 +541,19 @@ impl<EC: EvalCache> Program<EC> {
     /// Applies a custom transformation to the main term, assuming that it has been parsed but not
     /// yet transformed.
     ///
-    /// The term is left in whatever state it started in.
-    ///
-    /// This state-management isn't great, as it breaks the usual linear order of state changes.
-    /// In particular, there's no protection against double-applying the same transformation, and no
-    /// protection against applying it to a term that's in an unexpected state.
-    pub fn custom_transform<E, F>(&mut self, mut transform: F) -> Result<(), CacheError<E>>
+    /// If multiple invocations of `custom_transform` are needed, each subsequent invocation must supply
+    /// `transform_id` with with a number higher than that of all previous invocations.
+    pub fn custom_transform<E, F>(
+        &mut self,
+        transform_id: usize,
+        mut transform: F,
+    ) -> Result<(), CacheError<E>>
     where
         F: FnMut(&mut CacheHub, RichTerm) -> Result<RichTerm, E>,
     {
         self.vm
             .import_resolver_mut()
-            .custom_transform(self.main_id, &mut transform)
+            .custom_transform(self.main_id, transform_id, &mut transform)
     }
 
     /// Retrieve the parsed term, typecheck it, and generate a fresh initial environment. If

diff --git a/lsp/vscode-extension/yarn.lock b/lsp/vscode-extension/yarn.lock
@@ -2522,9 +2522,9 @@ table@^6.0.9:
     strip-ansi "^6.0.1"
 
 tar-fs@^2.0.0:
-  version "2.1.3"
-  resolved "https://registry.yarnpkg.com/tar-fs/-/tar-fs-2.1.3.tgz#fb3b8843a26b6f13a08e606f7922875eb1fbbf92"
-  integrity sha512-090nwYJDmlhwFwEW3QQl+vaNnxsO2yVsd45eTKRBzSzu+hlb1w2K9inVq5b0ngXuLVqQ4ApvsUHHnu/zQNkWAg==
+  version "2.1.4"
+  resolved "https://registry.yarnpkg.com/tar-fs/-/tar-fs-2.1.4.tgz#800824dbf4ef06ded9afea4acafe71c67c76b930"
+  integrity sha512-mDAjwmZdh7LTT6pNleZ05Yt65HC3E+NiQzl672vQG38jIrehtJk/J3mNwIg+vShQPcLF/LV7CMnDW6vjj6sfYQ==
   dependencies:
     chownr "^1.1.1"
     mkdirp-classic "^0.5.2"