From 3e88b36a406d52d4aa44bd23f1aef3dc431e396f Mon Sep 17 00:00:00 2001 From: Taus Date: Wed, 26 Aug 2026 22:05:42 +0000 Subject: [PATCH] yeast: Support optional guards in rules Permits the use of guards of the form `where expr` (with `expr` being any Rust expression that evaluates to a boolean) inside of rules. The guard should come after the query itself, and before the `=>` that separates the query from the body, like so: ``` (foo bar: _? @bar) where bar.is_some() => (baz bar: {bar}) ``` If the guard is absent, it behaves as if it were `where true`. Inside of the guard body, all captures are treated as if they are raw. Translation of non-raw captures only happens if the guard succeeds. The guard can also access the user-defined context. This context is shared with the body if the guard succeeds, and discarded if the guard fails (just as it is currently discarded after running a rule body). (I think it's unlikely that we'll ever want to mutate the context from inside the guard, but you never know...) --- shared/yeast-macros/src/lib.rs | 18 +++ shared/yeast-macros/src/parse.rs | 141 +++++++++++++++------- shared/yeast/doc/yeast.md | 48 ++++++++ shared/yeast/src/lib.rs | 166 ++++++++++++++----------- shared/yeast/tests/test.rs | 201 +++++++++++++++++++++++++++++++ 5 files changed, 459 insertions(+), 115 deletions(-) diff --git a/shared/yeast-macros/src/lib.rs b/shared/yeast-macros/src/lib.rs index 82b0e3e7b408..e9f7b46d0570 100644 --- a/shared/yeast-macros/src/lib.rs +++ b/shared/yeast-macros/src/lib.rs @@ -122,6 +122,15 @@ pub fn trees(input: TokenStream) -> TokenStream { /// (output_template field: {name} {repeated}) /// ) /// +/// // A guard filters a successful query match. Every capture is raw in the +/// // guard because guards run before capture translation. +/// rule!( +/// (query_pattern field: _? @value) +/// where value.is_none() +/// => +/// (output_template) +/// ) +/// /// // Shorthand: captures become fields on the output node /// rule!((query ...) => output_kind) /// ``` @@ -131,6 +140,15 @@ pub fn trees(input: TokenStream) -> TokenStream { /// - `@name` (after `*`/`+`) → `name: Vec` /// - `@name` (after `?`) → `name: Option` /// +/// A guard is an optional Rust condition between the query and `=>`. It is +/// evaluated after the query matches and before any `@` captures are +/// translated. Returning false makes the driver try the next rule. All +/// captures are therefore raw in the guard; `@` versus `@@` controls only +/// whether the capture is translated for the transform. `ctx` provides +/// mutable user-context access, and `ast` provides read-only AST access. +/// Mutations to `ctx` are visible to the transform when the guard succeeds. +/// Omitting the guard is equivalent to writing `where true`. +/// /// `tree!` and `trees!` can be used without explicit context inside `{...}`. #[proc_macro] pub fn rule(input: TokenStream) -> TokenStream { diff --git a/shared/yeast-macros/src/parse.rs b/shared/yeast-macros/src/parse.rs index 34e859c912e9..87fb5435bcd1 100644 --- a/shared/yeast-macros/src/parse.rs +++ b/shared/yeast-macros/src/parse.rs @@ -645,6 +645,32 @@ enum CaptureMultiplicity { Repeated, } +fn capture_bindings<'a>(captures: impl Iterator) -> Vec { + captures + .map(|cap| { + let name = Ident::new(&cap.name, Span::call_site()); + let name_str = &cap.name; + match cap.multiplicity { + CaptureMultiplicity::Repeated => { + quote! { + let #name: Vec = __captures.get_all(#name_str); + } + } + CaptureMultiplicity::Optional => { + quote! { + let #name: Option = __captures.get_opt(#name_str); + } + } + CaptureMultiplicity::Single => { + quote! { + let #name: yeast::Id = __captures.get_var(#name_str).unwrap(); + } + } + } + }) + .collect() +} + /// Walk a token stream and extract all `@name` captures, noting whether /// they appear after `*` or `+` (repeated) or not. fn extract_captures(stream: &TokenStream) -> Vec { @@ -806,7 +832,7 @@ fn try_consume_return_annotation(tokens: &mut Tokens) -> Result Result { let mut tokens = input.into_iter().peekable(); - // Collect query tokens up to `=>` + // Collect query and optional `where` guard tokens up to `=>`. let mut query_tokens = Vec::new(); loop { match tokens.peek() { @@ -830,7 +856,7 @@ pub fn parse_rule_top(input: TokenStream) -> Result { } } - let query_stream: TokenStream = query_tokens.into_iter().collect(); + let (query_stream, guard) = split_rule_guard(query_tokens)?; // Extract captures from query let captures = extract_captures(&query_stream); @@ -838,40 +864,22 @@ pub fn parse_rule_top(input: TokenStream) -> Result { // Parse query let query_code = parse_query_top(query_stream.clone())?; + let (raw_captures, translated_captures): (Vec<_>, Vec<_>) = + captures.iter().partition(|capture| capture.raw); + // Capture names marked `@@name` (raw) — passed to the auto-translate // prefix as a skip list so those captures keep their input-schema ids. - let raw_capture_names: Vec<&str> = captures + let raw_capture_names: Vec<&str> = raw_captures .iter() - .filter(|c| c.raw) - .map(|c| c.name.as_str()) + .map(|capture| capture.name.as_str()) .collect(); - // Generate capture bindings + // Both capture sets are bound raw in the guard, which runs before + // auto-translation. In the transform, raw captures remain unchanged + // while translated captures are bound after auto-translation. let ctx_ident = Ident::new(IMPLICIT_CTX, Span::call_site()); - let bindings: Vec = captures - .iter() - .map(|cap| { - let name = Ident::new(&cap.name, Span::call_site()); - let name_str = &cap.name; - match cap.multiplicity { - CaptureMultiplicity::Repeated => { - quote! { - let #name: Vec = __captures.get_all(#name_str); - } - } - CaptureMultiplicity::Optional => { - quote! { - let #name: Option = __captures.get_opt(#name_str); - } - } - CaptureMultiplicity::Single => { - quote! { - let #name: yeast::Id = __captures.get_var(#name_str).unwrap(); - } - } - } - }) - .collect(); + let raw_bindings = capture_bindings(raw_captures.into_iter()); + let translated_bindings = capture_bindings(translated_captures.into_iter()); // Parse transform: the token(s) after `=>` fall into one of three // shapes, dispatched in order: @@ -1015,24 +1023,36 @@ pub fn parse_rule_top(input: TokenStream) -> Result { } }; + let guard = guard.unwrap_or_else(|| syn::parse_quote!(true)); Ok(quote! { { let __query = #query_code; - yeast::Rule::new(__query, Box::new(|__ast: &mut yeast::Ast, mut __captures: yeast::captures::Captures, __fresh: &yeast::tree_builder::FreshScope, __source_range: Option, __user_ctx: &mut _, __translator: yeast::TranslatorHandle<'_, _>| { - // Auto-translation prefix: recursively translate every - // captured node before invoking the user's transform body, - // except for `@@name` captures listed in `__skip` which the - // body consumes raw. - // For OneShot rules this preserves the legacy behaviour - // (input-schema captures translated to output-schema - // nodes); for Repeating rules it is a no-op. - let __skip: &[&str] = &[#(#raw_capture_names),*]; - __translator.auto_translate_captures(&mut __captures, __ast, __user_ctx, __skip)?; - #(#bindings)* - let mut #ctx_ident = yeast::build::BuildCtx::with_translator(__ast, &__captures, __fresh, __source_range, __user_ctx, __translator); - let __result: Vec = { #transform_body }; - Ok(__result) - })) + yeast::Rule::guarded( + __query, + Box::new(|__ast: &yeast::Ast, __captures: &yeast::captures::Captures, __user_ctx: &mut _| { + #(#raw_bindings)* + #(#translated_bindings)* + let ast = __ast; + let #ctx_ident = __user_ctx; + Ok(#guard) + }), + Box::new(|__ast: &mut yeast::Ast, mut __captures: yeast::captures::Captures, __fresh: &yeast::tree_builder::FreshScope, __source_range: Option, __user_ctx: &mut _, __translator: yeast::TranslatorHandle<'_, _>| { + // Auto-translation prefix: recursively translate every + // captured node before invoking the user's transform body, + // except for `@@name` captures listed in `__skip` which the + // body consumes raw. + // For OneShot rules this preserves the legacy behaviour + // (input-schema captures translated to output-schema + // nodes); for Repeating rules it is a no-op. + let __skip: &[&str] = &[#(#raw_capture_names),*]; + __translator.auto_translate_captures(&mut __captures, __ast, __user_ctx, __skip)?; + #(#raw_bindings)* + #(#translated_bindings)* + let mut #ctx_ident = yeast::build::BuildCtx::with_translator(__ast, &__captures, __fresh, __source_range, __user_ctx, __translator); + let __result: Vec = { #transform_body }; + Ok(__result) + }), + ) } }) } @@ -1041,6 +1061,37 @@ pub fn parse_rule_top(input: TokenStream) -> Result { // Token utilities // --------------------------------------------------------------------------- +/// Split the tokens before a rule's `=>` into its query and optional +/// top-level `where guard`. Groups are opaque `TokenTree`s, so a `where` +/// inside the query or guard expression is not mistaken for the separator. +fn split_rule_guard(tokens: Vec) -> Result<(TokenStream, Option)> { + let guard_index = tokens + .iter() + .position(|tok| matches!(tok, TokenTree::Ident(ident) if ident == "where")); + + let Some(guard_index) = guard_index else { + return Ok((tokens.into_iter().collect(), None)); + }; + + let query: TokenStream = tokens[..guard_index].iter().cloned().collect(); + if query.is_empty() { + return Err(syn::Error::new( + Span::call_site(), + "expected query before rule guard", + )); + } + + let guard_tokens: TokenStream = tokens[guard_index + 1..].iter().cloned().collect(); + if guard_tokens.is_empty() { + return Err(syn::Error::new_spanned( + tokens[guard_index].clone(), + "expected expression after rule guard `where`", + )); + } + let guard = syn::parse2::(guard_tokens)?; + Ok((query, Some(guard))) +} + fn peek_is_at(tokens: &mut Tokens) -> bool { matches!(tokens.peek(), Some(TokenTree::Punct(p)) if p.as_char() == '@') } diff --git a/shared/yeast/doc/yeast.md b/shared/yeast/doc/yeast.md index be3bc913c700..3e3e1cd3610f 100644 --- a/shared/yeast/doc/yeast.md +++ b/shared/yeast/doc/yeast.md @@ -458,6 +458,54 @@ yeast::rule!( The shorthand `=> kind` form auto-generates the template, mapping each capture name to a field of the same name on the output node. +### Guards + +A rule may include a Rust guard between its query and `=>`. The guard runs +after the query matches but before any captures are translated. If it returns +`false`, the rule is treated as a non-match and the driver tries the next rule. +Omitting the guard is equivalent to writing `where true`: + +```rust +rule!( + (tupleExpr + elements: (labeledExpr + label: _? @label + expression: @inner) + elements: _* @rest) + where label.is_none() && rest.is_empty() + => + expr { inner } +) +``` + +Every capture is a raw input-schema id in the guard, regardless of whether it +uses `@` or `@@`, because the guard runs before translation. The marker controls +the transform binding only: `@inner` is translated after the guard accepts the +rule, while a capture marked `@@` would remain raw in the transform as well. +In this example `label` and `rest` are empty whenever the guard succeeds, so +there is nothing to translate for those captures. + +Guards receive the mutable user context as `ctx` and the raw AST as `ast`. +The framework clones the user context before evaluating each guard. If the +guard succeeds, its context mutations are visible to the rule transform and +recursive translation; if it fails, the clone is discarded before the next +rule is tried: + +```rust +rule!( + (tupleExpr elements: _* @elements) + where { + ctx.in_pattern = true; + elements + .first() + .and_then(|element| ast.get_node(*element)) + .is_some() + } + => + (tuple_pattern element: {elements}) +) +``` + ### Annotation form Rules that need imperative logic — mutating [`BuildCtx`] state per diff --git a/shared/yeast/src/lib.rs b/shared/yeast/src/lib.rs index c290246541f5..94f75faa4763 100644 --- a/shared/yeast/src/lib.rs +++ b/shared/yeast/src/lib.rs @@ -1014,8 +1014,16 @@ pub type Transform = Box< + Sync, >; +/// Predicate evaluated after a rule's query matches and before its transform +/// runs. Returning `false` makes the rule behave as if its query did not +/// match, so the driver continues to the next rule. The user context is a +/// private clone for this rule attempt; mutations are retained for the +/// transform when the guard succeeds and discarded when it fails. +pub type Guard = Box Result + Send + Sync>; + pub struct Rule { query: QueryNode, + guard: Option>, transform: Transform, /// If true, after this rule fires on a node the engine will try to /// re-apply this same rule on the result root. Defaults to false: @@ -1025,9 +1033,22 @@ pub struct Rule { } impl Rule { + /// Construct an unguarded rule from its query and transform. pub fn new(query: QueryNode, transform: Transform) -> Self { Self { query, + guard: None, + transform, + repeated: false, + } + } + + /// Construct a guarded rule. The guard sees raw captures and shares its + /// private user-context clone with the transform when it succeeds. + pub fn guarded(query: QueryNode, guard: Guard, transform: Transform) -> Self { + Self { + query, + guard: Some(guard), transform, repeated: false, } @@ -1042,30 +1063,29 @@ impl Rule { self } - fn try_rule( - &self, - ast: &mut Ast, - node: Id, - fresh: &tree_builder::FreshScope, - user_ctx: &mut C, - translator: TranslatorHandle<'_, C>, - ) -> Result>, String> { - match self.try_match(ast, node)? { - Some(captures) => Ok(Some( - self.run_transform(ast, captures, node, fresh, user_ctx, translator)?, - )), - None => Ok(None), + /// Attempt to match this rule's query against `node`, returning the raw + /// captures on success. Does not evaluate the guard or invoke the + /// transform. + fn match_query(&self, ast: &Ast, node: Id) -> Result, String> { + let mut captures = Captures::new(); + if !self.query.do_match(ast, node, &mut captures)? { + return Ok(None); } + Ok(Some(captures)) } - /// Attempt to match this rule's query against `node`, returning the - /// resulting captures on success. Does not invoke the transform. - fn try_match(&self, ast: &Ast, node: Id) -> Result, String> { - let mut captures = Captures::new(); - if self.query.do_match(ast, node, &mut captures)? { - Ok(Some(captures)) + /// Evaluate this rule's guard against a successful query match. An + /// unguarded rule always succeeds. + fn guard_matches( + &self, + ast: &Ast, + captures: &Captures, + user_ctx: &mut C, + ) -> Result { + if let Some(guard) = &self.guard { + guard(ast, captures, user_ctx) } else { - Ok(None) + Ok(true) } } @@ -1154,13 +1174,19 @@ fn apply_repeating_rules_inner( if Some(rule_ptr) == skip_rule { continue; } - // 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 Some(captures) = rule.match_query(ast, id)? else { + continue; + }; + + // Give each structurally-matching rule a private clone of the user + // context before its guard runs. Guard mutations are visible to the + // transform and recursive translation when the guard succeeds, but a + // failed guard drops the clone before trying the next rule. let mut local = user_ctx.clone(); + if !rule.guard_matches(ast, &captures, &mut local)? { + continue; + } + // 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 @@ -1168,29 +1194,25 @@ fn apply_repeating_rules_inner( let translator = TranslatorHandle { inner: TranslatorImpl::Repeating, }; - 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 - // query doesn't loop. Other rules and child traversal are - // unaffected. - let next_skip = if rule.repeated { None } else { Some(rule_ptr) }; - let mut results = Vec::new(); - for node in result_node { - results.extend(apply_repeating_rules_inner( - index, - ast, - &mut local, - node, - fresh, - rewrite_depth + 1, - next_skip, - )?); - } - return Ok(results); + let result_nodes = rule.run_transform(ast, captures, id, fresh, &mut local, translator)?; + + // For non-repeated rules, suppress further application of *this* + // rule on the result root, so a rule whose output matches its own + // query doesn't loop. Other rules and child traversal are unaffected. + let next_skip = if rule.repeated { None } else { Some(rule_ptr) }; + let mut results = Vec::new(); + for node in result_nodes { + results.extend(apply_repeating_rules_inner( + index, + ast, + &mut local, + node, + fresh, + rewrite_depth + 1, + next_skip, + )?); } - // Rule didn't match; `local` is dropped as we loop to the next - // rule. + return Ok(results); } // Take the parent's fields by ownership: the recursion will rewrite @@ -1271,29 +1293,33 @@ fn apply_one_shot_rules_inner( let node_kind = ast.get_node(id).map(|n| n.kind_name()).unwrap_or(""); for rule in index.rules_for_kind(node_kind) { - if let Some(captures) = rule.try_match(ast, id)? { - // 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 - // capture up front, preserving the legacy behavior). - let translator = TranslatorHandle { - inner: TranslatorImpl::OneShot { - index, - fresh, - rewrite_depth, - matched_root: id, - }, - }; - let result = rule.run_transform(ast, captures, id, fresh, &mut local, translator)?; - return Ok(result); + let Some(captures) = rule.match_query(ast, id)? else { + continue; + }; + + // Clone only after the query matches, but before the guard runs. + // Guard mutations are visible to the transform and transitively- + // translated captures when the guard succeeds; a failed guard drops + // the clone before trying the next rule. + let mut local = user_ctx.clone(); + if !rule.guard_matches(ast, &captures, &mut local)? { + continue; } + + // 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 capture up front, + // preserving the legacy behavior). + let translator = TranslatorHandle { + inner: TranslatorImpl::OneShot { + index, + fresh, + rewrite_depth, + matched_root: id, + }, + }; + let result = rule.run_transform(ast, captures, id, fresh, &mut local, translator)?; + return Ok(result); } Err(format!( diff --git a/shared/yeast/tests/test.rs b/shared/yeast/tests/test.rs index bdc17c7593dc..60ed8afe84dd 100644 --- a/shared/yeast/tests/test.rs +++ b/shared/yeast/tests/test.rs @@ -918,6 +918,104 @@ fn test_shorthand_rule() { ); } +#[derive(Clone, Default)] +struct GuardTestContext { + enabled: bool, + selected: bool, +} + +fn guarded_integer_rules() -> Vec> { + vec![ + yeast::rule!( + (integer) @@raw + where { + ctx.selected = ctx.enabled && ast.source_text(raw) == "1"; + ctx.selected + } + => + identifier { + assert!(ctx.selected); + tree!((identifier "guarded")) + } + ), + yeast::rule!((integer) => (identifier "fallback")), + ] +} + +#[test] +fn test_rule_guard_reads_raw_capture_and_user_context() { + fn run(input: &str, enabled: bool) -> String { + let lang: tree_sitter::Language = tree_sitter_ruby::LANGUAGE.into(); + let schema = + yeast::node_types_yaml::schema_from_yaml_with_language(OUTPUT_SCHEMA_YAML, &lang) + .unwrap(); + let phases = vec![Phase::new( + "test", + PhaseKind::Repeating, + guarded_integer_rules(), + )]; + let runner = Runner::with_schema(lang, &schema, &phases); + let mut user_ctx = GuardTestContext { + enabled, + selected: false, + }; + let ast = runner.run_with_ctx(input, &mut user_ctx).unwrap(); + dump_ast(&ast, ast.get_root(), input) + } + + assert_dump_eq( + &run("1", true), + r#" + program + identifier "guarded" + "#, + ); + assert_dump_eq( + &run("2", true), + r#" + program + identifier "fallback" + "#, + ); + assert_dump_eq( + &run("1", false), + r#" + program + identifier "fallback" + "#, + ); +} + +#[test] +fn test_rule_guard_binds_optional_and_repeated_raw_captures() { + fn rules() -> Vec { + vec![ + yeast::rule!( + (array (_)? @@first (_)* @@rest) + where first.is_some() && rest.is_empty() + => + (identifier "single") + ), + yeast::rule!((array) => (identifier "multiple")), + ] + } + + assert_dump_eq( + &run_and_dump("[1]", rules()), + r#" + program + identifier "single" + "#, + ); + assert_dump_eq( + &run_and_dump("[1, 2]", rules()), + r#" + program + identifier "multiple" + "#, + ); +} + #[test] fn test_chained_rules_output_only_kind() { // Exercise rule chaining where an intermediate kind exists only in the @@ -1134,6 +1232,84 @@ fn test_one_shot_phase_errors_when_no_rule_matches() { ); } +#[test] +fn test_one_shot_guard_runs_before_capture_translation() { + let lang: tree_sitter::Language = tree_sitter_ruby::LANGUAGE.into(); + let schema = + yeast::node_types_yaml::schema_from_yaml_with_language(OUTPUT_SCHEMA_YAML, &lang).unwrap(); + let rules: Vec = vec![ + yeast::rule!( + (program (_)* @stmts) + => + (program stmt: {stmts}) + ), + // There is deliberately no OneShot rule for `identifier`. If this + // rule translated `@left` before binding it raw in the guard, the run + // would fail instead of evaluating the guard and falling through to + // the next assignment rule. + yeast::rule!( + (assignment left: (_) @left) + where ast.source_text(left) == "not-x" + => + (first_node left: {left}) + ), + yeast::rule!((assignment) => (identifier "fallback")), + ]; + let phases = vec![Phase::new("translate", PhaseKind::OneShot, rules)]; + let runner: Runner = Runner::with_schema(lang, &schema, &phases); + + let input = "x = 1"; + let ast = runner.run(input).unwrap(); + let dump = dump_ast(&ast, ast.get_root(), input); + assert_dump_eq( + &dump, + r#" + program + stmt: identifier "fallback" + "#, + ); +} + +#[test] +fn test_one_shot_guard_context_mutation_is_visible_to_transform() { + let lang: tree_sitter::Language = tree_sitter_ruby::LANGUAGE.into(); + let schema = + yeast::node_types_yaml::schema_from_yaml_with_language(OUTPUT_SCHEMA_YAML, &lang).unwrap(); + let rules: Vec> = vec![ + yeast::rule!( + (program (_)* @stmts) + => + (program stmt: {stmts}) + ), + yeast::rule!( + (integer) + where { + ctx.selected = true; + true + } + => + identifier { + assert!(ctx.selected); + tree!((identifier "guarded")) + } + ), + ]; + let phases = vec![Phase::new("translate", PhaseKind::OneShot, rules)]; + let runner = Runner::with_schema(lang, &schema, &phases); + let mut user_ctx = GuardTestContext::default(); + + let input = "1"; + let ast = runner.run_with_ctx(input, &mut user_ctx).unwrap(); + let dump = dump_ast(&ast, ast.get_root(), input); + assert_dump_eq( + &dump, + r#" + program + stmt: identifier "guarded" + "#, + ); +} + /// OneShot recursion must apply rules to *captured* nodes, even if the rule /// returns a captured child verbatim. A buggy implementation that only /// recurses into the children of the rule's output (rather than into the @@ -1549,6 +1725,31 @@ fn test_rules_macro_accepts_bare_shorthand_form() { ); } +#[test] +fn test_rules_macro_accepts_bare_guarded_rule() { + let rules: Vec = yeast::rules! { + input: "tests/input-types.yml", + output: "tests/node-types.yml", + [ + (integer) @@raw + where ast.source_text(raw) == "1" + => + (identifier "guarded"), + + (integer) => (identifier "fallback"), + ] + }; + + let dump = run_and_dump("1", rules); + assert_dump_eq( + &dump, + r#" + program + identifier "guarded" + "#, + ); +} + /// Backwards-compat: explicit `rule!(...)` invocations inside `rules!` /// should still type-check and behave the same as the bare form. #[test]