fix commit feedback

Author: gouslu

rust/otap-dataflow/crates/engine-macros/src/capability.rs

@@ -28,9 +28,18 @@
 //! These limitations are fundamental to the type-erased `HashMap<TypeId, Entry>`
 //! registry design and the `SharedAsLocal` adapter delegation pattern:
 //!
-//! - **Trait-level generics / lifetimes** — `TypeId::of::<T>()` requires a
-//!   concrete monomorphized type; a generic registration struct cannot produce
-//!   a single `TypeId` for the registry.
+//! - **Trait-level generics or lifetime parameters** — e.g.
+//!   `#[capability] trait Foo<T>` or `#[capability] trait Bar<'a>`.
+//!   Method-level generics (one bullet up) *are* supported because the
+//!   trait object `dyn Foo` exists for the trait as a whole. Trait-level
+//!   parameters, by contrast, mean the trait isn't one type but a
+//!   *family* — `Foo<u32>`, `Foo<String>`, etc. — each with its own
+//!   `TypeId::of::<Foo<T>>()`. The registry keys entries by a single
+//!   `TypeId` per capability, so there is no monomorphized type for the
+//!   generated registration struct to advertise. Concretize the
+//!   parameter at the trait definition (e.g. work over `String` or a
+//!   sealed enum) or split the family into separate `#[capability]`
+//!   traits.
 //! - **Supertraits** (`trait Foo: Bar`) — the `SharedAsLocal` adapter only
 //!   delegates methods defined directly on the `#[capability]` trait. It cannot
 //!   auto-implement supertrait methods. Define all methods directly on the

rust/otap-dataflow/crates/engine/src/capability/registry/capabilities.rs

@@ -48,14 +48,25 @@ impl Capabilities {
     ///
     /// # One-shot contract
     ///
-    /// Each capability binding may be claimed **at most once per
-    /// node** across all four accessors (`require_local`,
-    /// `require_shared`, `optional_local`, `optional_shared`). Node
-    /// factories are expected to call this once at construction,
-    /// store the returned handle, and clone/share it within the node
-    /// as needed. A second call for the same capability — including
-    /// the implicit shared claim made by the `SharedAsLocal` fallback
-    /// — returns [`Error::CapabilityAlreadyConsumed`].
+    /// Each resolved entry is one-shot per node: a `require_local`
+    /// claim consumes the local entry, and a `require_shared` claim
+    /// consumes the shared entry. Node factories are expected to
+    /// call each accessor at most once at construction, store the
+    /// returned handle, and clone/share it within the node as needed.
+    ///
+    /// **SharedAsLocal fallback.** When a binding's extension only
+    /// registered a shared variant, `require_local` falls through to
+    /// the shared entry and routes its output through the
+    /// capability's adapter. In that case there is only one
+    /// underlying produce closure, so `require_local` and
+    /// `require_shared` share a single one-shot guard \u2014 claiming
+    /// either accessor consumes the binding for both. A second call
+    /// returns [`Error::CapabilityAlreadyConsumed`].
+    ///
+    /// **Native dual.** When a binding's extension registered both a
+    /// native local and a native shared variant, the two are
+    /// distinct objects with independent guards: a node may claim
+    /// `require_local` and `require_shared` and receive both.
     ///
     /// # Errors
     ///
@@ -75,25 +86,54 @@ impl Capabilities {
         &self,
     ) -> Result<Rc<C::Local>, Error> {
         let id = TypeId::of::<C>();
+
+        // Native local path. `Cell::take()` is the one-shot guard.
+        if let Some(entry) = self.local.get(&id) {
+            let produce = entry
+                .produce
+                .take()
+                .ok_or_else(|| Error::CapabilityAlreadyConsumed {
+                    capability: C::name().to_owned(),
+                })?;
+            let rc_any = produce();
+            let trait_object = rc_any
+                .downcast_ref::<Rc<C::Local>>()
+                .cloned()
+                .unwrap_or_else(|| {
+                    panic!(
+                        "BUG: capability '{}': local entry type mismatch in registry",
+                        C::name(),
+                    )
+                });
+            entry.tracker_consumed.set(true);
+            return Ok(trait_object);
+        }
+
+        // SharedAsLocal fallback. The same `Cell::take()` on the
+        // shared entry is the binding's one-shot guard, so claiming
+        // the local-via-fallback accessor here naturally consumes
+        // the native shared accessor too — a subsequent
+        // `require_shared` returns [`Error::CapabilityAlreadyConsumed`].
         let entry = self
-            .local
+            .shared
             .get(&id)
             .ok_or_else(|| Error::CapabilityNotBound {
                 capability: C::name().to_owned(),
                 execution_model: "local",
             })?;
-        if entry.claimed.replace(true) {
-            return Err(Error::CapabilityAlreadyConsumed {
+        let produce = entry
+            .produce
+            .take()
+            .ok_or_else(|| Error::CapabilityAlreadyConsumed {
                 capability: C::name().to_owned(),
-            });
-        }
-        let rc_any = (entry.produce)();
+            })?;
+        let rc_any = (entry.adapt_as_local)(produce());
         let trait_object = rc_any
             .downcast_ref::<Rc<C::Local>>()
             .cloned()
             .unwrap_or_else(|| {
                 panic!(
-                    "BUG: capability '{}': local entry type mismatch in registry",
+                    "BUG: capability '{}': SharedAsLocal adapter type mismatch in registry",
                     C::name(),
                 )
             });
@@ -135,12 +175,13 @@ impl Capabilities {
                 capability: C::name().to_owned(),
                 execution_model: "shared",
             })?;
-        if entry.claimed.replace(true) {
-            return Err(Error::CapabilityAlreadyConsumed {
+        let produce = entry
+            .produce
+            .take()
+            .ok_or_else(|| Error::CapabilityAlreadyConsumed {
                 capability: C::name().to_owned(),
-            });
-        }
-        let trait_object = (entry.produce)()
+            })?;
+        let trait_object = produce()
             .downcast::<Box<C::Shared>>()
             .map(|b| *b)
             .unwrap_or_else(|_| {
@@ -174,7 +215,9 @@ impl Capabilities {
         &self,
     ) -> Result<Option<Rc<C::Local>>, Error> {
         let id = TypeId::of::<C>();
-        if !self.local.contains_key(&id) {
+        // Available either as a native local entry or as a
+        // SharedAsLocal fallback through the shared entry.
+        if !self.local.contains_key(&id) && !self.shared.contains_key(&id) {
             return Ok(None);
         }
         self.require_local::<C>().map(Some)

rust/otap-dataflow/crates/engine/src/capability/registry/entry.rs

@@ -141,31 +141,57 @@ impl SharedCapabilityEntry {
 // ── Resolved entries (per-node) ─────────────────────────────────────────────
 
 /// A resolved local capability entry for a specific node.
+///
+/// Inserted only for **native** local registrations. The
+/// SharedAsLocal fallback path does *not* produce a separate local
+/// entry — instead, [`Capabilities::require_local`] falls through to
+/// the [`ResolvedSharedEntry`] when no native local entry exists,
+/// taking the shared `produce` cell and routing its output through
+/// the entry's `adapt_as_local` fn pointer. This collapses the
+/// fallback's two former entries into one, so the per-binding
+/// one-shot guard is encoded directly in the shared entry's
+/// `Cell::take()` — no separate claim flag needed.
 pub(crate) struct ResolvedLocalEntry {
     /// Per-node produce closure, cloned from the registry entry.
-    pub(crate) produce: Box<dyn LocalProduce>,
-    /// Per-node one-shot consumption flag for the consumer-side
-    /// "claim each binding at most once per node" contract enforced
-    /// by [`Capabilities`](super::Capabilities).
-    pub(crate) claimed: Cell<bool>,
-    /// Cross-node consumption flag for the underlying extension
-    /// variant, shared with the [`ConsumedTracker`].
     ///
-    /// - For native local bindings: cell lives under the tracker's
-    ///   *local* bucket for the `(capability, extension)` pair.
-    /// - For `SharedAsLocal` bindings: cell lives under the tracker's
-    ///   *shared* bucket — consuming the adapter is considered a use
-    ///   of the shared variant.
+    /// Wrapped in `Cell<Option<_>>` so the consumer-side one-shot
+    /// guard is encoded in the type: `Some` = unclaimed, `None` =
+    /// already claimed. Consumption is a single `Cell::take()` that
+    /// atomically moves the closure out and marks the binding as
+    /// claimed; a second `take()` returns `None` and yields
+    /// [`Error::CapabilityAlreadyConsumed`](super::Error::CapabilityAlreadyConsumed).
+    pub(crate) produce: Cell<Option<Box<dyn LocalProduce>>>,
+    /// Cross-node consumption flag for the local variant, shared
+    /// with the [`ConsumedTracker`].
     pub(crate) tracker_consumed: Rc<Cell<bool>>,
 }
 
 /// A resolved shared capability entry for a specific node.
+///
+/// Serves two roles. As a native shared entry, [`Capabilities::require_shared`]
+/// takes its `produce` cell to mint a `Box<dyn C::Shared>`. As the
+/// SharedAsLocal fallback for a binding with no native local entry,
+/// [`Capabilities::require_local`] takes the same cell and routes the
+/// output through `adapt_as_local` to mint an `Rc<dyn C::Local>`.
+/// Either path consumes the single underlying [`Cell::take()`], so the
+/// per-binding one-shot contract is enforced naturally without an
+/// auxiliary claim flag.
 pub(crate) struct ResolvedSharedEntry {
     /// Per-node produce closure, cloned from the registry entry.
-    pub(crate) produce: Box<dyn SharedProduce>,
-    /// Per-node one-shot consumption flag.
-    pub(crate) claimed: Cell<bool>,
+    ///
+    /// Wrapped in `Cell<Option<_>>` for the same reason as
+    /// [`ResolvedLocalEntry::produce`]: `Cell::take()` doubles as
+    /// the one-shot consumer-side guard, and is shared between the
+    /// native shared and SharedAsLocal-fallback paths.
+    pub(crate) produce: Cell<Option<Box<dyn SharedProduce>>>,
     /// Cross-node consumption flag, shared with the
     /// [`ConsumedTracker`].
     pub(crate) tracker_consumed: Rc<Cell<bool>>,
+    /// Adapter fn pointer used by the SharedAsLocal fallback path
+    /// in [`Capabilities::require_local`]. Takes a produced shared
+    /// trait object (erased as `Box<dyn Any + Send>`) and returns a
+    /// local trait object (erased as `Rc<dyn Any>` wrapping
+    /// `Rc<dyn C::Local>`). Originally minted by the
+    /// `#[capability]`-generated `shared_entry::<E>` caster.
+    pub(crate) adapt_as_local: fn(Box<dyn Any + Send>) -> Rc<dyn Any>,
 }

rust/otap-dataflow/crates/engine/src/capability/registry/mod.rs

@@ -39,7 +39,7 @@ pub use storage::CapabilityRegistry;
 
 // ── Crate-internal re-exports ────────────────────────────────────────────────
 
-pub(crate) use entry::{LocalProduce, ResolvedLocalEntry, ResolvedSharedEntry};
+pub(crate) use entry::{ResolvedLocalEntry, ResolvedSharedEntry};
 // TODO(extension-system): wired by engine build phase in a follow-up PR;
 // until then only the test module imports it via `use super::*;`.
 #[allow(unused_imports)]

rust/otap-dataflow/crates/engine/src/capability/registry/resolve.rs

@@ -6,13 +6,12 @@
 //! [`Capabilities`] for consumption.
 
 use super::{
-    Capabilities, CapabilityRegistry, ConsumedTracker, Error, LocalProduce, ResolvedLocalEntry,
+    Capabilities, CapabilityRegistry, ConsumedTracker, Error, ResolvedLocalEntry,
     ResolvedSharedEntry,
 };
 use otap_df_config::ExtensionId;
 use std::any::TypeId;
 use std::collections::{HashMap, HashSet};
-use std::rc::Rc;
 
 /// Resolves a node's capability bindings against the registry.
 ///
@@ -108,13 +107,17 @@ pub(crate) fn resolve_bindings(
             )));
         }
 
-        // Resolve local entry: prefer a native local registration; else,
-        // if the extension registered a shared entry, invoke the
-        // capability's `SharedAsLocal` adapter to build a new local
-        // wrapper around this node's own clone of the shared instance.
-        let native_local = local_entry;
-
-        if let Some(local_entry) = native_local {
+        // Resolve local entry. Native local registrations get their
+        // own resolved entry; the SharedAsLocal fallback is *not*
+        // materialized here — instead, [`Capabilities::require_local`]
+        // falls through to the shared entry below and runs its
+        // `adapt_as_local` adapter at consumption time. This collapses
+        // the fallback's two former entries into one, so the
+        // per-binding one-shot guard is the shared entry's
+        // `Cell::take()`: claiming the local-via-fallback execution
+        // model naturally consumes the native shared execution model
+        // too.
+        if let Some(local_entry) = local_entry {
             let tracker_consumed = tracker.ensure_local_consumer_slot(
                 cap_type_id,
                 known_cap.name,
@@ -124,8 +127,7 @@ pub(crate) fn resolve_bindings(
             let prior = local_entries.insert(
                 cap_type_id,
                 ResolvedLocalEntry {
-                    produce: local_entry.produce.clone_box(),
-                    claimed: std::cell::Cell::new(false),
+                    produce: std::cell::Cell::new(Some(local_entry.produce.clone_box())),
                     tracker_consumed,
                 },
             );
@@ -134,56 +136,11 @@ pub(crate) fn resolve_bindings(
                 "resolve_bindings: duplicate local entry for capability '{cap_name_str}' \
                  - the config layer should prevent two bindings with the same capability name",
             );
-        } else if let Some(shared_entry) = shared_entry {
-            // SharedAsLocal fallback: defer the shared mint to the
-            // (single, one-shot) `require_local` call. The closure
-            // captures a refcounted handle to the shared produce
-            // closure plus the capability's `adapt_as_local` fn
-            // pointer, then on invocation produces a new shared
-            // instance and routes it through the adapter to yield a
-            // type-erased `Rc<dyn Any>` wrapping the local trait
-            // object.
-            let adapt = shared_entry.adapt_as_local;
-            let produce_shared: Rc<Box<dyn super::entry::SharedProduce>> =
-                Rc::new(shared_entry.produce.clone_box());
-            let local_produce: Box<dyn LocalProduce> = Box::new(move || adapt((*produce_shared)()));
-
-            // The extension only provided a shared variant — record
-            // the adapter consumer's cross-node consumption against
-            // the shared bucket. The `// Resolve shared entry` block
-            // below also calls `ensure_shared_consumer_slot` for
-            // this `(cap, ext)` pair; that method is get-or-insert,
-            // so both users share the same `Rc<Cell<bool>>` for
-            // tracker accounting. Per-node one-shot enforcement is
-            // tracked separately in each entry's `claimed` cell, so
-            // a node *can* claim both the local-via-shared and
-            // native-shared sides of a fallback binding.
-            let tracker_consumed = tracker.ensure_shared_consumer_slot(
-                cap_type_id,
-                known_cap.name,
-                shared_entry.extension_id.clone(),
-            );
-
-            let prior = local_entries.insert(
-                cap_type_id,
-                ResolvedLocalEntry {
-                    produce: local_produce,
-                    claimed: std::cell::Cell::new(false),
-                    tracker_consumed,
-                },
-            );
-            debug_assert!(
-                prior.is_none(),
-                "resolve_bindings: duplicate local entry for capability '{cap_name_str}' \
-                 (SharedAsLocal fallback)",
-            );
         }
 
-        // Resolve shared entry. If the SharedAsLocal fallback above
-        // already called `ensure_shared_consumer_slot` for this
-        // `(cap, ext)` pair, this call is a no-op lookup that returns
-        // the same cell — get-or-insert idempotency keeps the two
-        // users pointing at one slot.
+        // Resolve shared entry. Used both as the native shared
+        // binding and as the SharedAsLocal fallback target when no
+        // native local entry exists.
         if let Some(shared_entry) = shared_entry {
             let tracker_consumed = tracker.ensure_shared_consumer_slot(
                 cap_type_id,
@@ -194,9 +151,9 @@ pub(crate) fn resolve_bindings(
             let prior = shared_entries.insert(
                 cap_type_id,
                 ResolvedSharedEntry {
-                    produce: shared_entry.produce.clone_box(),
-                    claimed: std::cell::Cell::new(false),
+                    produce: std::cell::Cell::new(Some(shared_entry.produce.clone_box())),
                     tracker_consumed,
+                    adapt_as_local: shared_entry.adapt_as_local,
                 },
             );
             debug_assert!(