Streamline dispatch of Assign operators

lionel- · lionel- · commit f2eeae6db26b · 2026-07-17T13:18:43.000+02:00
diff --git a/crates/oak_semantic/src/builder/builder_nse.rs b/crates/oak_semantic/src/builder/builder_nse.rs
@@ -4,10 +4,8 @@ use aether_syntax::RCall;
 use aether_syntax::RSyntaxKind;
 use biome_rowan::AstNode;
 use biome_rowan::AstNodeList;
-use biome_rowan::AstPtr;
 use biome_rowan::AstSeparatedList;
 use biome_rowan::TextRange;
-use oak_core::range::RangedAstPtr;
 use oak_core::syntax_ext::AnyRSelectorExt;
 use oak_core::syntax_ext::RIdentifierExt;
 
@@ -20,6 +18,7 @@ use super::SemanticIndexBuilder;
 use super::SourcedFile;
 use crate::effects::AssignBinding;
 use crate::effects::CallContext;
+use crate::effects::EffectSite;
 use crate::effects::Effects;
 use crate::effects::EffectsHandlers;
 use crate::effects::ResolvedArgumentEffect;
@@ -179,21 +178,24 @@ impl<R: ImportsResolver> SemanticIndexBuilder<R> {
 
     /// Scan a binary operator for an assign effect (e.g. magrittr's `x %<>% f()`)
     pub(super) fn scan_operator_assign(&mut self, bin: &RBinaryExpression) {
-        let Some(binding) = self.resolve_operator_assign(bin) else {
+        let Some(bindings) = self.resolve_operator_assign(bin) else {
             return;
         };
-        self.record_binding(binding.name.clone());
-        self.call_resolutions
-            .entry(bin.syntax().text_trimmed_range())
-            .or_default()
-            .assign
-            .push(binding);
+        let range = bin.syntax().text_trimmed_range();
+        for binding in bindings {
+            self.record_binding(binding.name.clone());
+            self.call_resolutions
+                .entry(range)
+                .or_default()
+                .assign
+                .push(binding);
+        }
     }
 
     /// Recognize a binding operator (`x %<>% f()`, `x %<~% expr`, `x := expr`)
-    /// as an assign effect and build its binding, or `None` for any other binary
+    /// as an assign effect and build its bindings, or `None` for any other binary
     /// operator.
-    fn resolve_operator_assign(&mut self, bin: &RBinaryExpression) -> Option<AssignBinding> {
+    fn resolve_operator_assign(&mut self, bin: &RBinaryExpression) -> Option<Vec<AssignBinding>> {
         let op = bin.operator().ok()?;
 
         // A binding operator is either a `%...%` (`SPECIAL`, e.g. `%<>%`, where
@@ -211,17 +213,8 @@ impl<R: ImportsResolver> SemanticIndexBuilder<R> {
         }
 
         let handlers = self.resolve_symbol_effects(op_text, bin.syntax().text_trimmed_range())?;
-        let _assign = handlers.assign?;
-
-        let left = bin.left().ok()?;
-        let right = bin.right().ok()?;
-        let (name, _) = assignment_name(&left)?;
-
-        Some(AssignBinding {
-            name,
-            name_expr: RangedAstPtr::new(&left),
-            value_expr: Some(AstPtr::new(&right)),
-        })
+        let ctx = CallContext::new();
+        handlers.assign?.resolve(EffectSite::Operator(bin), &ctx)
     }
 
     /// Copy the names a `Current + Lazy` body defines into the owner's
@@ -333,7 +326,7 @@ impl<R: ImportsResolver> SemanticIndexBuilder<R> {
             .and_then(|handler| handler.resolve(call, &ctx));
         let assign = handlers
             .assign
-            .and_then(|handler| handler.resolve(call, &ctx));
+            .and_then(|handler| handler.resolve(EffectSite::Call(call), &ctx));
 
         Some(Effects {
             arguments,
diff --git a/crates/oak_semantic/src/effects.rs b/crates/oak_semantic/src/effects.rs
@@ -2,12 +2,13 @@ use aether_syntax::AnyRArgumentName;
 use aether_syntax::AnyRExpression;
 use aether_syntax::AnyRValue;
 use aether_syntax::RArgument;
+use aether_syntax::RBinaryExpression;
 use aether_syntax::RCall;
 use biome_rowan::AstNode;
 use biome_rowan::AstPtr;
 use biome_rowan::AstSeparatedList;
 use biome_rowan::WalkEvent;
-// Re-exported so consumers building an `AssignBinding` (custom `EffectHandler`s)
+// Re-exported so consumers building an `AssignBinding` (custom `AssignHandler`s)
 // can name the `name_expr` field's type without depending on oak_core directly.
 pub use oak_core::range::RangedAstPtr;
 use oak_core::syntax_ext::RIdentifierExt;
@@ -53,7 +54,7 @@ pub struct EffectsHandlers {
     pub arguments: Option<&'static dyn EffectHandler<Output = ResolvedArgumentEffects>>,
     pub attach: Option<&'static dyn EffectHandler<Output = String>>,
     pub source: Option<&'static dyn EffectHandler<Output = Vec<String>>>,
-    pub assign: Option<&'static dyn EffectHandler<Output = Vec<AssignBinding>>>,
+    pub assign: Option<&'static dyn AssignHandler>,
 }
 
 /// Resolver for an effect of a call.
@@ -73,6 +74,25 @@ pub trait EffectHandler: std::fmt::Debug + Sync {
     fn resolve(&self, call: &RCall, ctx: &CallContext) -> Option<Self::Output>;
 }
 
+/// Where an effect is invoked. Most effects are only ever calls but an Assign
+/// effect can also be a binding operator (`x %<>% f`). [`AssignHandler`] takes
+/// this to disambiguate rather than a bare call.
+pub enum EffectSite<'a> {
+    Call(&'a RCall),
+    Operator(&'a RBinaryExpression),
+}
+
+/// Resolver for an assign-like effect.
+///
+/// Separate from [`EffectHandler`] because an assign has two invocation shapes,
+/// a call (`assign("x", v)`) and a binding operator (`x %<>% f`).
+///
+/// Contributed statically like [`EffectHandler`], so it's `Sync` for the
+/// registry `static`s.
+pub trait AssignHandler: std::fmt::Debug + Sync {
+    fn resolve(&self, site: EffectSite, ctx: &CallContext) -> Option<Vec<AssignBinding>>;
+}
+
 /// Context for effect handlers.
 ///
 /// Allows querying the properties or static values of arguments. Stateless
@@ -143,6 +163,15 @@ impl CallContext {
         }
     }
 
+    /// Read a quoted name argument. E.g. the LHS of an Assign operator.
+    pub fn resolve_quoted_symbol_or_string(&self, value: &AnyRExpression) -> Option<String> {
+        match value {
+            AnyRExpression::RIdentifier(ident) => Some(ident.name_text()),
+            AnyRExpression::AnyRValue(AnyRValue::RStringValue(s)) => s.string_text(),
+            _ => None,
+        }
+    }
+
     /// Statically evaluate an argument's value expression to a bool.
     pub fn resolve_static_bool(&self, value: &AnyRExpression) -> Option<bool> {
         match value {
@@ -454,10 +483,11 @@ pub struct AssignAnnotation {
     pub position: usize,
 }
 
-impl EffectHandler for AssignAnnotation {
-    type Output = Vec<AssignBinding>;
-
-    fn resolve(&self, call: &RCall, ctx: &CallContext) -> Option<Vec<AssignBinding>> {
+impl AssignHandler for AssignAnnotation {
+    fn resolve(&self, site: EffectSite, ctx: &CallContext) -> Option<Vec<AssignBinding>> {
+        let EffectSite::Call(call) = site else {
+            return None;
+        };
         let args = call.arguments().ok()?;
 
         // Matched positionally among unnamed arguments, same as `source`, so a
@@ -512,6 +542,30 @@ impl EffectHandler for AssignAnnotation {
     }
 }
 
+/// Handler for a binding operator (`x %<>% f()`, `x %<~% expr`, `x := expr`).
+///
+/// The operator captures its LHS unevaluated.
+#[derive(Debug, Clone, Copy)]
+pub struct BindingOperatorHandler;
+
+impl AssignHandler for BindingOperatorHandler {
+    fn resolve(&self, site: EffectSite, ctx: &CallContext) -> Option<Vec<AssignBinding>> {
+        let EffectSite::Operator(bin) = site else {
+            return None;
+        };
+        let left = bin.left().ok()?;
+        let right = bin.right().ok()?;
+
+        let name = ctx.resolve_quoted_symbol_or_string(&left)?;
+
+        Some(vec![AssignBinding {
+            name,
+            name_expr: RangedAstPtr::new(&left),
+            value_expr: Some(AstPtr::new(&right)),
+        }])
+    }
+}
+
 /// Match a named argument against `formals`. Returns the index of the matched
 /// formal.
 ///
diff --git a/crates/oak_semantic/src/effects_registry.rs b/crates/oak_semantic/src/effects_registry.rs
@@ -3,6 +3,7 @@ use crate::effects::ArgumentEffect;
 use crate::effects::ArgumentsAnnotation;
 use crate::effects::AssignAnnotation;
 use crate::effects::AttachAnnotation;
+use crate::effects::BindingOperatorHandler;
 use crate::effects::BquoteHandler;
 use crate::effects::EffectsHandlers;
 use crate::effects::SourceAnnotation;
@@ -121,7 +122,8 @@ macro_rules! source {
 }
 
 /// An assign entry: `(name-argument position)`. The function binds a name in the
-/// current scope.
+/// current scope, naming it in a positional argument it evaluates (`assign("x",
+/// v)`).
 macro_rules! assign {
     ($pkg:literal, $func:literal, $pos:literal) => {
         Entry {
@@ -137,6 +139,24 @@ macro_rules! assign {
     };
 }
 
+/// An assign-operator entry: a binding operator (`x %<>% f`, `x := v`) that binds
+/// a name in the current scope. It captures its LHS unevaluated, so the name
+/// comes from the LHS text rather than a positional argument, hence no position.
+macro_rules! assign_op {
+    ($pkg:literal, $func:literal) => {
+        Entry {
+            package: $pkg,
+            function: $func,
+            effects: EffectsHandlers {
+                arguments: None,
+                attach: None,
+                source: None,
+                assign: Some(&BindingOperatorHandler),
+            },
+        }
+    };
+}
+
 static REGISTRY: &[Entry] = &[
     // base NSE
     nse!("base", "evalq", ("expr", 0, Current, Eager)),
@@ -168,9 +188,9 @@ static REGISTRY: &[Entry] = &[
     assign!("base", "assign", 0),
     assign!("base", "delayedAssign", 0),
     // magrittr / rlang / S7 binding operators
-    assign!("magrittr", "%<>%", 0),
-    assign!("rlang", "%<~%", 0),
-    assign!("S7", ":=", 0),
+    assign_op!("magrittr", "%<>%"),
+    assign_op!("rlang", "%<~%"),
+    assign_op!("S7", ":="),
     // rlang
     nse!("rlang", "on_load", ("expr", 0, Current, Lazy)),
     // shiny
diff --git a/crates/oak_semantic/tests/integration/builder.rs b/crates/oak_semantic/tests/integration/builder.rs
@@ -6,8 +6,10 @@ use biome_rowan::AstNode;
 use biome_rowan::AstSeparatedList;
 use oak_semantic::build_index;
 use oak_semantic::effects::AssignBinding;
+use oak_semantic::effects::AssignHandler;
 use oak_semantic::effects::CallContext;
 use oak_semantic::effects::EffectHandler;
+use oak_semantic::effects::EffectSite;
 use oak_semantic::effects::RangedAstPtr;
 use oak_semantic::effects::SourceAnnotation;
 use oak_semantic::effects_registry;
@@ -2329,10 +2331,11 @@ struct MultiAssignHandler;
 
 static MULTI_ASSIGN_HANDLER: MultiAssignHandler = MultiAssignHandler;
 
-impl EffectHandler for MultiAssignHandler {
-    type Output = Vec<AssignBinding>;
-
-    fn resolve(&self, call: &RCall, _ctx: &CallContext) -> Option<Vec<AssignBinding>> {
+impl AssignHandler for MultiAssignHandler {
+    fn resolve(&self, site: EffectSite, _ctx: &CallContext) -> Option<Vec<AssignBinding>> {
+        let EffectSite::Call(call) = site else {
+            return None;
+        };
         // Point every binding's handles at the first argument. This test only
         // checks that multiple defs are created and resolve, not their ranges.
         let expr = call
