Fixing implementation of streams

Signed-off-by: Michael X. Grey <mxgrey@intrinsic.ai>

Author: mxgrey

src/builder.rs

@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2023 Open Source Robotics Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+*/
+
+use bevy::prelude::{Entity, Commands};
+
+use crate::{
+    Provider, UnusedTarget, StreamPack, Node, InputSlot, Output, StreamTargetMap,
+};
+
+pub(crate) mod internal;
+pub(crate) use internal::*;
+
+/// Device used for building a workflow. Simply pass a mutable borrow of this
+/// into any functions which ask for it.
+///
+/// Note that each scope has its own [`Builder`], and a panic will occur if a
+/// [`Builder`] gets used in the wrong scope. As of right now there is no known
+/// way to trick the compiler into using a [`Builder`] in the wrong scope, but
+/// please open an issue with a minimal reproducible example if you find a way
+/// to make it panic.
+pub struct Builder<'w, 's, 'a> {
+    pub(crate) scope: Entity,
+    pub(crate) commands: &'a mut Commands<'w, 's>,
+}
+
+impl<'w, 's, 'a> Builder<'w, 's, 'a> {
+    pub fn create_node<P: Provider>(
+        &mut self,
+        provider: P,
+    ) -> Node<P::Request, P::Response, P::Streams>
+    where
+        P::Request: 'static + Send + Sync,
+        P::Response: 'static + Send + Sync,
+        P::Streams: StreamPack,
+    {
+        let source = self.commands.spawn(()).id();
+        let target = self.commands.spawn(UnusedTarget).id();
+        provider.connect(source, target, self.commands);
+
+        let mut map = StreamTargetMap::default();
+        let (bundle, streams) = <P::Streams as StreamPack>::spawn_node_streams(
+            &mut map, self,
+        );
+        self.commands.entity(source).insert(bundle);
+        Node {
+            input: InputSlot::new(self.scope, source),
+            output: Output::new(self.scope, target),
+            streams,
+        }
+    }
+
+    /// Get the scope that this builder is building for
+    pub fn scope(&self) -> Entity {
+        self.scope
+    }
+
+    /// Borrow the commands for the builder
+    pub fn commands(&'a mut self) -> &'a mut Commands<'w, 's> {
+        &mut self.commands
+    }
+}

src/builder/internal.rs

@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 Open Source Robotics Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+*/
+
+use bevy::{
+    prelude::Entity,
+    ecs::system::Command,
+};
+
+use crate::{SingleTargetStorage, ForkTargetStorage};
+
+/// If two nodes have been created, they will each have a unique source and a
+/// target entity allocated to them. If we want to connect them, then we want
+/// the target of one to no longer be unique - we instead want it to be the
+/// source entity of the other. This [`Command`] redirects the target information
+/// of the sending node to target the source entity of the receiving node.
+struct RedirectConnection {
+    original_target: Entity,
+    new_target: Entity,
+}
+

src/chain.rs

@@ -17,17 +17,17 @@
 
 use std::future::Future;
 
+use bevy::prelude::Entity;
+
 use crate::{
     UnusedTarget, AddOperation, Node, InputSlot, Builder,
     ForkClone, StreamPack, Provider, ProvideOnce,
     AsMap, IntoBlockingMap, IntoAsyncMap, Output, Noop,
-    ForkTargetStorage,
+    ForkTargetStorage, StreamTargetMap,
     make_result_branching, make_cancel_filter_on_err,
     make_option_branching, make_cancel_filter_on_none,
 };
 
-use bevy::prelude::{Entity, Commands};
-
 use smallvec::SmallVec;
 
 pub mod fork_clone_builder;
@@ -76,6 +76,11 @@ impl<'w, 's, 'a, 'b, T: 'static + Send + Sync> Chain<'w, 's, 'a, 'b, T> {
         Output::new(self.builder.scope, self.target)
     }
 
+    /// Connect this output into an input slot.
+    pub fn connect(self, input: InputSlot<T>) {
+        TODO: Finish implementing this
+    }
+
     /// Connect the response at the end of the chain into a new provider. Get
     /// the response of the new provider as a chain so you can continue chaining
     /// operations.
@@ -105,10 +110,12 @@ impl<'w, 's, 'a, 'b, T: 'static + Send + Sync> Chain<'w, 's, 'a, 'b, T> {
         let source = self.target;
         let target = self.builder.commands.spawn(UnusedTarget).id();
         provider.connect(source, target, self.builder.commands);
+
+        let mut map = StreamTargetMap::default();
         let (bundle, streams) = <P::Streams as StreamPack>::spawn_node_streams(
-            self.builder.scope, self.builder.commands,
+            &mut map, self.builder,
         );
-        self.builder.commands.entity(source).insert(bundle);
+        self.builder.commands.entity(source).insert((bundle, map));
         Node {
             input: InputSlot::new(self.builder.scope, source),
             output: Output::new(self.builder.scope, target),

src/impulse.rs

@@ -23,7 +23,7 @@ use std::future::Future;
 
 use crate::{
     Promise, ProvideOnce, StreamPack, IntoBlockingMapOnce, IntoAsyncMapOnce,
-    AsMapOnce, UnusedTarget,
+    AsMapOnce, UnusedTarget, StreamTargetMap,
 };
 
 mod detach;
@@ -92,8 +92,9 @@ where
             self.target,
             TakenResponse::<Response>::new(response_sender),
         ));
-        let (bundle, stream_receivers) = Streams::take_streams(self.target, self.commands);
-        self.commands.entity(self.source).insert(bundle);
+        let mut map = StreamTargetMap::default();
+        let (bundle, stream_receivers) = Streams::take_streams(self.target, &mut map, self.commands);
+        self.commands.entity(self.source).insert((bundle, map));
 
         Recipient {
             response: response_promise,
@@ -203,21 +204,23 @@ where
             Store::<Response>::new(target),
         ));
 
+        let mut map = StreamTargetMap::default();
         let stream_targets = Streams::collect_streams(
-            self.source, target, self.commands,
+            self.source, target, &mut map, self.commands,
         );
-        self.commands.entity(self.source).insert(stream_targets);
+        self.commands.entity(self.source).insert((stream_targets, map));
     }
 
     /// Collect the stream data into [`Collection<T>`] components in the
     /// specified target, one collection for each stream data type. You must
     /// still decide what to do with the final response data.
     #[must_use]
     pub fn collect_streams(self, target: Entity) -> Impulse<'w, 's, 'a, Response, ()> {
+        let mut map = StreamTargetMap::default();
         let stream_targets = Streams::collect_streams(
-            self.source, target, self.commands,
+            self.source, target, &mut map, self.commands,
         );
-        self.commands.entity(self.source).insert(stream_targets);
+        self.commands.entity(self.source).insert((stream_targets, map));
 
         Impulse {
             source: self.source,
@@ -241,10 +244,11 @@ where
             Push::<Response>::new(target, false),
         ));
 
+        let mut map = StreamTargetMap::default();
         let stream_targets = Streams::collect_streams(
-            self.source, target, self.commands,
+            self.source, target, &mut map, self.commands,
         );
-        self.commands.entity(self.source).insert(stream_targets);
+        self.commands.entity(self.source).insert((stream_targets, map));
     }
 
     // TODO(@mxgrey): Consider offering ways for users to respond to cancellations.

src/lib.rs

@@ -15,6 +15,9 @@
  *
 */
 
+pub mod builder;
+pub use builder::*;
+
 pub mod cancel;
 pub use cancel::*;
 

src/operation/scope.rs

@@ -23,7 +23,7 @@ use crate::{
     BufferSettings, Buffer, Cancellable, OperationRoster, ManageCancellation,
     OperationError, OperationCancel, Cancel, UnhandledErrors, check_reachability,
     Blocker, Stream, StreamTargetStorage, StreamRequest, AddOperation,
-    ScopeSettings,
+    ScopeSettings, StreamTargetMap,
 };
 
 use backtrace::Backtrace;
@@ -1072,37 +1072,40 @@ pub(crate) struct ExitTarget {
 }
 
 pub(crate) struct RedirectScopeStream<T: Stream> {
+    target: Entity,
     _ignore: std::marker::PhantomData<T>,
 }
 
 impl<T: Stream> RedirectScopeStream<T> {
-    pub(crate) fn new() -> Self {
-        Self { _ignore: Default::default() }
+    pub(crate) fn new(target: Entity) -> Self {
+        Self { target, _ignore: Default::default() }
     }
 }
 
 impl<T: Stream> Operation for RedirectScopeStream<T> {
     fn setup(self, OperationSetup { source, world }: OperationSetup) -> OperationResult {
-        world.entity_mut(source).insert(
+        world.entity_mut(source).insert((
             InputBundle::<T>::new(),
-        );
+            SingleTargetStorage::new(self.target),
+        ));
         Ok(())
     }
 
     fn execute(
         OperationRequest { source, world, roster }: OperationRequest,
     ) -> OperationResult {
         let mut source_mut = world.get_entity_mut(source).or_broken()?;
+
+        // The target is optional because we want to "send" this stream even if
+        // there is no target listening, because streams may have custom sending
+        // behavior
+        let target = source_mut.get::<SingleTargetStorage>().map(|t| t.get());
         let Input { session: scoped_session, data } = source_mut.take_input::<T>()?;
-        let scope = source_mut.get::<ScopeStorage>().or_broken()?.get();
-        let stream_target = world
-            .get::<StreamTargetStorage<T>>(scope)
-            .map(|target| target.get());
         let parent_session = world.get::<ParentSession>(scoped_session).or_broken()?.get();
         data.send(StreamRequest {
             source,
             session: parent_session,
-            target: stream_target,
+            target,
             world,
             roster
         })
@@ -1161,8 +1164,10 @@ impl<T: Stream> Operation for RedirectWorkflowStream<T> {
         let exit_source = exit.source;
         let parent_session = exit.parent_session;
 
+        let stream_target_map = world.get::<StreamTargetMap>(exit_source).or_broken()?;
         let stream_target = world.get::<StreamTargetStorage<T>>(exit_source)
-            .map(|target| target.get());
+            .map(|target| stream_target_map.get(target.get()))
+            .flatten();
 
         data.send(StreamRequest {
             source,