Merge pull request #22136 from github/tausbn/yeast-translate-with-context

yeast: Add context scoping mechanism

Author: tausbn

shared/yeast/src/build.rs

@@ -161,52 +161,75 @@ impl<'a, C> BuildCtx<'a, C> {
 }
 
 impl<C: Clone> BuildCtx<'_, C> {
-    /// Recursively translate a node via the framework's rule machinery.
-    /// In a OneShot phase, applies OneShot rules to the given node and
-    /// returns the resulting node ids. In a Repeating phase, errors
-    /// (translation is not meaningful when input and output share a
-    /// schema).
+    /// Recursively translate every id in the given iterable via the
+    /// framework's rule machinery. In a OneShot phase, applies OneShot
+    /// rules to each id and returns the accumulated resulting node ids
+    /// in order. In a Repeating phase, errors (translation is not
+    /// meaningful when input and output share a schema).
+    ///
+    /// The single-`Id` case works too, because `Id: IntoIterator<Item
+    /// = Id>` is a singleton iterator — so `ctx.translate(some_id)?`
+    /// returns a `Vec<Id>` containing whatever `some_id` translated to.
     ///
     /// Errors if this `BuildCtx` was constructed by hand (without a
     /// translator handle) — for example, in unit tests that don't go
     /// through the rule driver.
-    pub fn translate<I: Into<Id>>(&mut self, id: I) -> Result<Vec<Id>, String> {
-        let id = id.into();
-        match &self.translator {
-            Some(t) => t.translate(self.ast, self.user_ctx, id),
-            None => Err("translate() called on a BuildCtx without a translator handle".into()),
+    pub fn translate<I: Into<Id>>(
+        &mut self,
+        ids: impl IntoIterator<Item = I>,
+    ) -> Result<Vec<Id>, String> {
+        let translator = self
+            .translator
+            .as_ref()
+            .ok_or("translate() called on a BuildCtx without a translator handle")?;
+        let mut out = Vec::new();
+        for id in ids {
+            let translated = translator.translate(self.ast, self.user_ctx, id.into())?;
+            out.extend(translated);
         }
+        Ok(out)
     }
 
-    /// Translate every node in an iterator with a **fresh** user context
-    /// (reset to `C::default()`), restoring the previous context afterwards.
+    /// Run `f` with a temporary child [`BuildCtx`] whose `user_ctx` is
+    /// a fresh clone of the current one, sharing everything else
+    /// (`ast`, `captures`, `fresh`, `source_range`, `translator`) by
+    /// re-borrow. Any mutations `f` makes to the child's `user_ctx`
+    /// are discarded when it returns — no restore needed, because the
+    /// mutations only ever happened on a local clone.
     ///
-    /// Use when descending into a subtree — a body, expression, or statement
-    /// list — that must not inherit any of the surrounding translation
-    /// context (for example an enclosing binding modifier). Accepts optional
-    /// (`Option<Id>`) and repeated (`Vec<Id>`) captures (both `IntoIterator`);
-    /// for a single `Id`, wrap it in `std::iter::once(id)`.
-    pub fn translate_reset<I: Into<Id>>(
-        &mut self,
-        ids: impl IntoIterator<Item = I>,
-    ) -> Result<Vec<Id>, String>
+    /// Use for the rare rule that needs to translate a subtree under a
+    /// modified context *and then continue using its own (unmodified)
+    /// context afterwards*. For rules where the modified translation
+    /// is the last use of `ctx`, mutate `ctx` in place — the framework
+    /// invokes each rule with a private clone of the user context, so
+    /// mutations are discarded on rule exit anyway.
+    ///
+    /// Example: an outer rule that translates one child subtree with a
+    /// reset context, then continues with the outer context intact:
+    ///
+    /// ```ignore
+    /// let val = ctx.scoped(|ctx| {
+    ///     ctx.reset();
+    ///     ctx.translate(val)
+    /// })?;
+    /// // `ctx` here is untouched by the reset inside the closure.
+    /// let other = ctx.translate(other_id)?;
+    /// ```
+    pub fn scoped<F, R>(&mut self, f: F) -> R
     where
-        C: Default,
+        F: for<'b> FnOnce(&mut BuildCtx<'b, C>) -> R,
     {
-        let saved = std::mem::take(&mut *self.user_ctx);
-        let mut out = Vec::new();
-        let mut result = Ok(());
-        for id in ids {
-            match self.translate(id) {
-                Ok(v) => out.extend(v),
-                Err(e) => {
-                    result = Err(e);
-                    break;
-                }
-            }
-        }
-        *self.user_ctx = saved;
-        result.map(|()| out)
+        let mut child_user_ctx = self.user_ctx.clone();
+        let mut child = BuildCtx {
+            ast: &mut *self.ast,
+            captures: self.captures,
+            fresh: self.fresh,
+            source_range: self.source_range,
+            user_ctx: &mut child_user_ctx,
+            translator: self.translator,
+        };
+        f(&mut child)
+        // child_user_ctx dropped; the outer `self` is unaffected.
     }
 }
 

shared/yeast/src/lib.rs

@@ -26,6 +26,12 @@ use query::QueryNode;
 /// without colliding with the impls for plain integers.
 ///
 /// Use `id.0` (or `id.into()`) to obtain the raw arena index.
+///
+/// Implements [`IntoIterator`] as a singleton (`iter::once(self)`)
+/// so that a bare `Id` can be used interchangeably with `Option<Id>`
+/// / `Vec<Id>` in places that expect an iterable of ids (e.g.
+/// [`crate::build::BuildCtx::translate`] and the field-splice
+/// interpolation via [`IntoFieldIds`]).
 #[repr(transparent)]
 #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug, Hash, Serialize)]
 pub struct Id(pub usize);
@@ -42,16 +48,25 @@ impl From<Id> for usize {
     }
 }
 
+impl IntoIterator for Id {
+    type Item = Id;
+    type IntoIter = std::iter::Once<Id>;
+    fn into_iter(self) -> Self::IntoIter {
+        std::iter::once(self)
+    }
+}
+
 /// Field and Kind ids are provided by tree-sitter
 type FieldId = u16;
 type KindId = u16;
 
 /// Trait for values that can be appended to a field's id list inside a
 /// `tree!`/`trees!`/`rule!` template (in `{expr}` placeholders).
 ///
-/// `Id` pushes a single id; the blanket impl for
-/// `IntoIterator<Item: Into<Id>>` handles `Vec<Id>`, `Option<Id>`,
-/// arbitrary iterators yielding `Id`, etc.
+/// The blanket impl for `IntoIterator<Item: Into<Id>>` handles all
+/// current shapes: `Vec<Id>`, `Option<Id>`, arbitrary iterators
+/// yielding `Id`, and a bare `Id` itself (which is `IntoIterator`
+/// via a singleton).
 ///
 /// This lets `{expr}` interpolate any of these shapes without a
 /// dedicated splice syntax — the macro emits the same trait-dispatched
@@ -60,12 +75,6 @@ pub trait IntoFieldIds {
     fn extend_into(self, out: &mut Vec<Id>);
 }
 
-impl IntoFieldIds for Id {
-    fn extend_into(self, out: &mut Vec<Id>) {
-        out.push(self);
-    }
-}
-
 impl<I, T> IntoFieldIds for I
 where
     I: IntoIterator<Item = T>,
@@ -732,6 +741,16 @@ pub struct TranslatorHandle<'a, C> {
     inner: TranslatorImpl<'a, C>,
 }
 
+// Manual `Copy` / `Clone` so `TranslatorHandle<'_, C>: Copy` holds
+// regardless of whether `C: Copy`. `TranslatorImpl` contains only
+// shared references, which are `Copy` unconditionally.
+impl<C> Copy for TranslatorHandle<'_, C> {}
+impl<C> Clone for TranslatorHandle<'_, C> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
 /// Internal phase-specific translation state. Kept private — callers
 /// interact with [`TranslatorHandle`] only.
 enum TranslatorImpl<'a, C> {
@@ -752,6 +771,16 @@ enum TranslatorImpl<'a, C> {
     Repeating,
 }
 
+// Manual `Copy` / `Clone` so `TranslatorImpl<'_, C>: Copy` holds
+// regardless of whether `C: Copy`. All variants hold only shared
+// references and small `Copy` scalars.
+impl<C> Copy for TranslatorImpl<'_, C> {}
+impl<C> Clone for TranslatorImpl<'_, C> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
 impl<'a, C: Clone> TranslatorHandle<'a, C> {
     /// Recursively apply OneShot rules to `id` and return the resulting
     /// node ids. Errors in a Repeating phase (where translation is not
@@ -970,18 +999,21 @@ fn apply_repeating_rules_inner<C: Clone>(
         if Some(rule_ptr) == skip_rule {
             continue;
         }
-        // Snapshot the user context before invoking the rule so that any
-        // mutations the rule makes are visible during recursive translation
-        // of its result, but not leaked to the parent's siblings.
-        let snapshot = user_ctx.clone();
+        // Give each rule attempt a private clone of the user context.
+        // Any mutations the rule makes are visible to its transform and
+        // to the recursive translation of its result, but never leak
+        // back to the parent — the clone is simply dropped when we
+        // return. This is also `?`-safe: an error return drops `local`
+        // without touching the caller's `user_ctx`.
+        let mut local = user_ctx.clone();
         // Repeating rules don't need a real translator: their captures
         // aren't auto-translated (Repeating preserves the input schema),
         // and `ctx.translate(id)` errors if invoked from a Repeating
         // transform.
         let translator = TranslatorHandle {
             inner: TranslatorImpl::Repeating,
         };
-        let try_result = rule.try_rule(ast, id, fresh, user_ctx, translator)?;
+        let try_result = rule.try_rule(ast, id, fresh, &mut local, translator)?;
         if let Some(result_node) = try_result {
             // For non-repeated rules, suppress further application of *this*
             // rule on the result root, so a rule whose output matches its own
@@ -993,19 +1025,17 @@ fn apply_repeating_rules_inner<C: Clone>(
                 results.extend(apply_repeating_rules_inner(
                     index,
                     ast,
-                    user_ctx,
+                    &mut local,
                     node,
                     fresh,
                     rewrite_depth + 1,
                     next_skip,
                 )?);
             }
-            *user_ctx = snapshot;
             return Ok(results);
         }
-        // Rule didn't match; restore any speculative changes (none expected
-        // since try_rule only mutates on match, but be defensive).
-        *user_ctx = snapshot;
+        // Rule didn't match; `local` is dropped as we loop to the next
+        // rule.
     }
 
     // Take the parent's fields by ownership: the recursion will rewrite
@@ -1087,11 +1117,13 @@ fn apply_one_shot_rules_inner<C: Clone>(
 
     for rule in index.rules_for_kind(node_kind) {
         if let Some(captures) = rule.try_match(ast, id)? {
-            // Snapshot the user context before invoking the rule so that any
-            // mutations the rule (or its transitively-translated captures)
-            // make are visible during this rule's transform, but not leaked
-            // to the parent's siblings.
-            let snapshot = user_ctx.clone();
+            // Give the rule a private clone of the user context. Any
+            // mutations the rule (or its transitively-translated
+            // captures) make are visible during this rule's transform,
+            // but never leak back — the clone is dropped when we
+            // return. `?`-safe: an error return drops `local` without
+            // touching the caller's `user_ctx`.
+            let mut local = user_ctx.clone();
             // Build the translator handle the transform will use to
             // recursively translate captures (or, for macro-generated
             // rules, the auto-translate prefix uses it to translate every
@@ -1104,8 +1136,7 @@ fn apply_one_shot_rules_inner<C: Clone>(
                     matched_root: id,
                 },
             };
-            let result = rule.run_transform(ast, captures, id, fresh, user_ctx, translator)?;
-            *user_ctx = snapshot;
+            let result = rule.run_transform(ast, captures, id, fresh, &mut local, translator)?;
             return Ok(result);
         }
     }