Skip to content

yeast: Add context scoping mechanism #22136

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tausbn merged 4 commits into from
Jul 9, 2026
+159 −79
Merged

yeast: Add context scoping mechanism #22136

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8272848
yeast: Implement IntoIterator for Id as a singleton
tausbn Jul 4, 2026
13510b1
yeast: Add BuildCtx::scoped for isolated context modification
tausbn Jul 4, 2026
9b7a9c3
yeast: Use clone-and-drop for per-rule context isolation
tausbn Jul 6, 2026
33da3ef
yeast: Fix typo
tausbn Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 58 additions & 36 deletions shared/yeast/src/build.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,52 +161,74 @@ 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>` as a singleton iterator — so `ctx.translate(some_id)?`
Comment thread
tausbn marked this conversation as resolved.
Outdated
/// 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's rule-boundary save/restore cleans up on rule exit.
///
/// 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)
/// })?;
Comment thread
jketema marked this conversation as resolved.
/// // `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.
}
}

Expand Down
20 changes: 20 additions & 0 deletions shared/yeast/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -741,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> {
Expand All @@ -761,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
Expand Down
59 changes: 42 additions & 17 deletions unified/extractor/src/languages/swift/swift.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,29 @@ impl SwiftContext {
///
/// True exactly when an enclosing binding has published its modifier into
/// `outer_modifiers`. This is reliable because non-binding subtrees
/// (bodies, initializer values, ...) are translated with a reset context
/// (see `BuildCtx::translate_reset`), so a bare identifier only sees a
/// (bodies, initializer values, ...) are translated after resetting the
/// context (see `reset`), so a bare identifier only sees a
/// non-empty `outer_modifiers` when it really is a binding.
fn in_binding_pattern(&self) -> bool {
!self.outer_modifiers.is_empty()
}

/// Clear the context fields that must not propagate into an
/// expression / statement / body subtree.
///
/// Mirrors `Default::default()` for `SwiftContext` today, but is a
/// named method so future context fields can opt in or out of
/// clearing here per-field.
///
/// Called before recursively translating a body / initializer
/// slot. Most rules mutate `ctx` in place — the framework's
/// rule-boundary snapshot/restore cleans up on exit. Rules that
/// need the outer context intact *after* the reset-and-translate
/// (see e.g. the `property_binding` willSet/didSet rule) wrap the
/// mutation in `ctx.scoped(...)` instead.
fn reset(&mut self) {
*self = SwiftContext::default();
}
}

/// Build a freshly-created `chained_declaration` modifier node if
Expand Down Expand Up @@ -239,7 +256,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
name: (identifier #{name})
type: {ty}
accessor_kind: (accessor_kind "get")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// Stored property with willSet/didSet observers (initializer
// optional) → a `variable_declaration` followed by one
Expand All @@ -260,12 +277,20 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
observers: (willset_didset_block willset: _? @@ws didset: _? @@ds))
=>
{{
// The initializer value must not inherit the binding context
// (it may contain patterns, e.g. a switch expression), so
// translate it with a reset context. The observers keep the
// context: each willSet/didSet accessor emits the binding
// modifier and resets its own body.
let val = ctx.translate_reset(val)?;
// The initializer value must not inherit the binding
// context (it may contain patterns, e.g. a switch
// expression), so translate it inside a `ctx.scoped`
// block — the block receives a temporary `ctx` whose
// `user_ctx` is a clone; mutations to it are discarded
// when the block returns, so the outer `ctx` is intact
// for the observer loop below. The observers keep the
// outer context: each willSet/didSet accessor emits
// the binding modifier and, in turn, resets the
// context for its own body.
let val = ctx.scoped(|ctx| {
ctx.reset();
ctx.translate(val)
})?;
Comment on lines +291 to +294

@jketema jketema Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This begs the question: why do we need scoped here, but not in any of the other places where we had translate_reset before?

@tausbn tausbn Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The difference here is that we want to use ctx (or more specifically the user-controlled user_ctx field on it -- which through some Deref magic is directly accessible on ctx itself) further down in the body of this rule. When we get to ctx.translate(obs), we want that user context to not have been reset. But we've already put var_decl into the result, and this required var to have been translated.

Thus, we have to introduce a temporary context -- or at least it's the easiest way to get the desired behaviour. (I guess technically we could do the var stuff at the end, and then prepend it onto result but I think that's less clear than the current implementation.)

@jketema jketema Jul 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me rephrase, because that's not really helping me. I assume that a context here roughly corresponds to some type of scope in Swift. If so, from the perspective of the Swift AST, what is different here than in the other cases where we had translate_reset.

@tausbn tausbn Jul 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The context is a mechanism for passing information from a node down to its descendants. For instance, In some cases we have an outer type that we want to propagate down into some descendant. In this case, we set the appropriate part of the SwiftContext (which is just a struct that we can choose the content of), specifically the following field:

codeql/unified/extractor/src/languages/swift/swift.rs

Line 23 in 33da3ef

property_type: Option<yeast::Id>,

In the previous code, we had two different kinds of translate_reset calls. The simple uses just had things like

(accessor_declaration
  body: (block stmt: {ctx.translate_reset(body)?})
)

This had the effect of translating body within an empty context. (That is, it temporarily saves the current context, then creates an empty context, runs the translation, and then restores the saved context.)

But in this case, there's no reason to restore the context -- it gets dropped anyway afterwards, which is why we can just rewrite it to use reset and then the normal translate.

(accessor_declaration
  body: (block stmt: {ctx.reset(); ctx.translate(body)?})
)

The difficult case is the following (eliding the unimportant bits):

// previously: val = translate_reset(val)?;
let val = ctx.scoped(|ctx| {
    ctx.reset();
    ctx.translate(val)
})?;

let var_decl = tree!(
    (variable_declaration
        modifier: {ctx.outer_modifiers.clone()}
        modifier: {chained_modifier(&mut ctx)}
        pattern: (name_pattern identifier: (identifier #{name}))
        type: {ty}
        value: {val})
);

// Publish the property name for the observer rules.
ctx.property_name = Some(tree!((identifier #{name})));
// Observers are subsequent outputs of this flattening
// rule, so they always get `chained_declaration`.
ctx.is_chained = true;

let mut result = vec![var_decl];
for obs in ws.into_iter().chain(ds) {
    result.extend(ctx.translate(obs)?);
}
result

The behaviour is the same as before -- we just express it slightly differently. The important thing is that when we do ctx.translate(val) we need the context to be cleared, otherwise things get translated incorrectly. However, when we do ctx.translate(obs), it shouldn't be cleared (or at least, that's the current behaviour). So, currently we use translate_reset to provide a temporary empty context for the val translation, and for the rewrite we explicitly call ctx.scoped, which allows us to make any changes to ctx that we like without it affecting anything outside the scope.

@jketema jketema Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is still too abstract. What Swift construct are we trying to transform/translate here?

@tausbn tausbn Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's translating the property_binding node, propagating the name of the property to the willSet/didSet observers.

I don't see why this is relevant to the present PR, however. The behaviour of the rule is unchanged -- all I've done is change the yeast API to make it a bit more flexible and fine-grained.

Perhaps we should have a meeting where we go through what the rules are doing? (Though I should point out that I cannot speak to the motivation for representing things the way they are currently. For that you'll have to talk to asgerf.)

@jketema jketema Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see why this is relevant to the present PR, however.

I would like to understand whether this change is correct, but I find this very hard based only only the code and the abstract descriptions you have given.

@tausbn tausbn Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, let's break it down. The change uses ctx.scoped, so let's start there:

    pub fn scoped<F, R>(&mut self, f: F) -> R
    where
        F: for<'b> FnOnce(&mut BuildCtx<'b, C>) -> R,
    {
        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.
    }

We know what the function f does: it calls ctx.reset() followed by ctx.translate(val), so let's inline that. (Let's pretend that val is somehow magically bound inside this function.)

    pub fn scoped<F, R>(&mut self, f: F) -> R
    where
        F: for<'b> FnOnce(&mut BuildCtx<'b, C>) -> R,
    {
        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,
        };
        child.reset();
        child.translate(val)
        // child_user_ctx dropped; the outer `self` is unaffected.
    }

Now, child.reset() dispatches to child.user_ctx.reset(), which has the effect of setting user_ctx to its default value. Thus, the above is essentially

    pub fn scoped<F, R>(&mut self, f: F) -> R
    where
        F: for<'b> FnOnce(&mut BuildCtx<'b, C>) -> R,
    {
        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 BuildCtx::default(), // <-- empty context
            translator: self.translator,
        };
        child.translate(val)
        // child_user_ctx dropped; the outer `self` is unaffected.
    }

In other words, we call translate on val inside a context where user_ctx has been reset to its default state, and return the result of this translation. After that, the temporarily reset context goes away, our reborrowed values (like ast etc.) are handed back to the outer context, and execution continues in the outer rule.

Finally, translate (without _reset) just goes through the given Ids and translates them one by one:

    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)
    }

Now let's consider translate_reset:

    pub fn translate_reset<I: Into<Id>>(
        &mut self,
        ids: impl IntoIterator<Item = I>,
    ) -> Result<Vec<Id>, String>
    where
        C: Default,
    {
        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)
    }

Here, we first save user_ctx using std::mem::take. This has the effect of resetting the field self.user_ctx to its default value. Then we go through ids, translating each one inside the current context (which has user_ctx reset), and gather up the resulting mapped nodes in result. Finally, at the end we reassign user_ctx to its original saved value and return result.

(Now, you may wonder: can't user_ctx modifications leak through from one translate call to the next, as we go through ids? And the answer is "no" because translate calls apply_one_shot_rules_inner, and this function also makes a local clone of user_ctx which is the thing that is actually passed to the rewrite rule. Thus, each rewrite rule has its own isolated view of the user context. It's shared downwards to its descendants, but it can never leak upwards or sideways.)

For the cases where translate_reset is the last thing run in a rewrite rule block, that final act of restoring the user context is completely superfluous, since the user context gets dropped immediately anyway. With scoped the restoration is implicit -- inside the scope we're working with a copy of ctx that has a user_ctx that is freshly cloned, and which gets dropped at the end of the scope. Thus, those changes never leak out, and the original context is implicitly restored at the end of the scope.

Does this help? If not, we should pair on doing a deeper dive into the code, since I'm not sure I can explain it more clearly through the present medium.

@jketema jketema Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. This helps, but in an unexpected way: it confirms what I understood from your code changes, so that's nice ❤️.

I'm convinced at this point that the behavior is unchanged.


let var_decl = tree!(
(variable_declaration
Expand Down Expand Up @@ -295,7 +320,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
// The enclosing `property_declaration` leads `ctx.outer_modifiers`
// with the `let`/`var` binding modifier, so the auto-translated name
// pattern (the LHS) becomes a binding, while the initializer value is
// translated with a reset context (see `translate_reset`).
// translated with a reset context (see `SwiftContext::reset`).
rule!(
(property_binding
name: @pattern
Expand All @@ -307,7 +332,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
modifier: {chained_modifier(&mut ctx)}
pattern: {pattern}
type: {ty}
value: {ctx.translate_reset(val)?}) // reset context: the initializer must not see the binding
value: {ctx.reset(); ctx.translate(val)?})
),
// property_declaration: flatten declarators (each may translate
// to multiple nodes — variable_declaration and/or
Expand Down Expand Up @@ -1118,7 +1143,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
name: {ctx.property_name.ok_or("computed_getter outside property_binding context")?}
type: {ctx.property_type}
accessor_kind: (accessor_kind "get")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// Computed setter with explicit parameter name.
rule!(
Expand All @@ -1131,7 +1156,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
type: {ctx.property_type}
accessor_kind: (accessor_kind "set")
parameter: (parameter pattern: (name_pattern identifier: (identifier #{param})))
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// Computed setter without explicit parameter name; body optional.
rule!(
Expand All @@ -1143,7 +1168,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
name: {ctx.property_name.ok_or("computed_setter outside property_binding context")?}
type: {ctx.property_type}
accessor_kind: (accessor_kind "set")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// Computed modify → accessor_declaration
rule!(
Expand All @@ -1155,7 +1180,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
name: {ctx.property_name.ok_or("computed_modify outside property_binding context")?}
type: {ctx.property_type}
accessor_kind: (accessor_kind "modify")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// willset/didset block — spread to children (only reachable as a
// fallback; the outer property_binding manual rule normally
Expand All @@ -1173,7 +1198,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
modifier: {chained_modifier(&mut ctx)}
name: {ctx.property_name.ok_or("willset_clause outside property_binding context")?}
accessor_kind: (accessor_kind "willSet")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// didset clause → accessor_declaration (body optional).
rule!(
Expand All @@ -1184,7 +1209,7 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
modifier: {chained_modifier(&mut ctx)}
name: {ctx.property_name.ok_or("didset_clause outside property_binding context")?}
accessor_kind: (accessor_kind "didSet")
body: (block stmt: {ctx.translate_reset(body)?}))
body: (block stmt: {ctx.reset(); ctx.translate(body)?}))
),
// Preprocessor conditionals — unsupported
rule!((diagnostic) => (unsupported_node)),
Expand Down