Reimplementing Join using buffers

Signed-off-by: Michael X. Grey <grey@openrobotics.org>

Author: mxgrey

src/buffer.rs

@@ -17,6 +17,14 @@
 
 use bevy::prelude::Entity;
 
+use crate::{Builder, Chain, UnusedTarget, OnNewBufferValue, InputSlot};
+
+mod buffered;
+pub use buffered::*;
+
+mod bufferable;
+pub use bufferable::*;
+
 /// A buffer is a special type of node within a workflow that is able to store
 /// and release data. When a session is finished, the buffered data from the
 /// session will be automatically cleared.
@@ -26,6 +34,44 @@ pub struct Buffer<T> {
     pub(crate) _ignore: std::marker::PhantomData<T>,
 }
 
+impl<T> Buffer<T> {
+    /// Get a unit `()` trigger output each time a new value is added to the buffer.
+    pub fn on_new_value<'w, 's, 'a, 'b>(
+        &self,
+        builder: &'b mut Builder<'w, 's, 'a>
+    ) -> Chain<'w, 's, 'a, 'b, ()> {
+        assert_eq!(self.scope, builder.scope);
+        let target = builder.commands.spawn(UnusedTarget).id();
+        builder.commands.add(OnNewBufferValue::new(self.source, target));
+        Chain::new(target, builder)
+    }
+
+    /// Specify that you want to pull from this Buffer by cloning. This can be
+    /// used by operations like join to tell them that they should clone from
+    /// the buffer instead of consuming from it.
+    pub fn by_cloning(self) -> CloneFromBuffer<T>
+    where
+        T: Clone,
+    {
+        CloneFromBuffer {
+            scope: self.scope,
+            source: self.source,
+            _ignore: Default::default()
+        }
+    }
+
+    /// Get an input slot for this buffer.
+    pub fn input_slot(self) -> InputSlot<T> {
+        InputSlot::new(self.scope, self.source)
+    }
+}
+
+pub struct CloneFromBuffer<T: Clone> {
+    pub(crate) scope: Entity,
+    pub(crate) source: Entity,
+    pub(crate) _ignore: std::marker::PhantomData<T>,
+}
+
 /// Settings to describe the behavior of a buffer.
 #[derive(Default, Clone, Copy)]
 pub struct BufferSettings {
@@ -62,3 +108,27 @@ impl Default for RetentionPolicy {
         Self::KeepLast(1)
     }
 }
+
+impl<T> Clone for Buffer<T> {
+    fn clone(&self) -> Self {
+        Self {
+            scope: self.scope,
+            source: self.source,
+            _ignore: Default::default(),
+        }
+    }
+}
+
+impl<T> Copy for Buffer<T> {}
+
+impl<T: Clone> Clone for CloneFromBuffer<T> {
+    fn clone(&self) -> Self {
+        Self {
+            scope: self.scope,
+            source: self.source,
+            _ignore: Default::default(),
+        }
+    }
+}
+
+impl<T: Clone> Copy for CloneFromBuffer<T> {}

src/buffer/bufferable.rs

@@ -0,0 +1,111 @@
+/*
+ * 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 crate::{
+    Buffer, CloneFromBuffer, Output, Builder, BufferSettings, Buffered, Join,
+    UnusedTarget, AddOperation,
+};
+
+pub trait Bufferable {
+    type BufferType: Buffered;
+
+    /// Convert these bufferable workflow elements into buffers if they are not
+    /// buffers already.
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType;
+
+    /// Join these bufferable workflow elements.
+    fn join(
+        self,
+        builder: &mut Builder,
+    ) -> Output<<Self::BufferType as Buffered>::Item>
+    where
+        Self: Sized,
+        Self::BufferType: 'static + Send + Sync,
+        <Self::BufferType as Buffered>::Item: 'static + Send + Sync,
+    {
+        let buffers = self.as_buffer(builder);
+        let join = builder.commands.spawn(()).id();
+        let target = builder.commands.spawn(UnusedTarget).id();
+        builder.commands.add(AddOperation::new(join, Join::new(buffers, target)));
+
+        Output::new(builder.scope, target)
+    }
+}
+
+impl<T: 'static + Send + Sync> Bufferable for Buffer<T> {
+    type BufferType = Self;
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        assert_eq!(self.scope, builder.scope());
+        self
+    }
+}
+
+impl<T: 'static + Send + Sync + Clone> Bufferable for CloneFromBuffer<T> {
+    type BufferType = Self;
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        assert_eq!(self.scope, builder.scope());
+        self
+    }
+}
+
+impl<T: 'static + Send + Sync> Bufferable for Output<T> {
+    type BufferType = Buffer<T>;
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        assert_eq!(self.scope(), builder.scope());
+        let buffer = builder.create_buffer::<T>(BufferSettings::default());
+        builder.connect(self, buffer.input_slot());
+        buffer
+    }
+}
+
+impl<T0, T1> Bufferable for (T0, T1)
+where
+    T0: Bufferable,
+    T1: Bufferable,
+{
+    type BufferType = (T0::BufferType, T1::BufferType);
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        (
+            self.0.as_buffer(builder),
+            self.1.as_buffer(builder),
+        )
+    }
+}
+
+impl<T0, T1, T2> Bufferable for (T0, T1, T2)
+where
+    T0: Bufferable,
+    T1: Bufferable,
+    T2: Bufferable,
+{
+    type BufferType = (T0::BufferType, T1::BufferType, T2::BufferType);
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        (
+            self.0.as_buffer(builder),
+            self.1.as_buffer(builder),
+            self.2.as_buffer(builder),
+        )
+    }
+}
+
+
+impl<T: Bufferable, const N: usize> Bufferable for [T; N] {
+    type BufferType = [T::BufferType; N];
+    fn as_buffer(self, builder: &mut Builder) -> Self::BufferType {
+        self.map(|b| b.as_buffer(builder))
+    }
+}

src/buffer/buffered.rs

@@ -0,0 +1,196 @@
+/*
+ * 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, World};
+
+use smallvec::SmallVec;
+
+use crate::{
+    Buffer, CloneFromBuffer, OperationError, OrBroken, InspectInput, ManageInput,
+};
+
+pub trait Buffered: Copy {
+    fn buffered_count(
+        &self,
+        session: Entity,
+        world: &World,
+    ) -> Result<usize, OperationError>;
+
+    type Item;
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError>;
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]>;
+}
+
+impl<T: 'static + Send + Sync> Buffered for Buffer<T> {
+    fn buffered_count(&self, session: Entity, world: &World) -> Result<usize, OperationError> {
+        world.get_entity(self.source).or_broken()?
+            .buffered_count::<T>(session)
+    }
+
+    type Item = T;
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError> {
+        world.get_entity_mut(self.source).or_broken()?
+            .pull_from_buffer::<T>(session)
+    }
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]> {
+        SmallVec::from_iter([self.source])
+    }
+}
+
+impl<T: 'static + Send + Sync + Clone> Buffered for CloneFromBuffer<T> {
+    fn buffered_count(
+        &self,
+        session: Entity,
+        world: &World,
+    ) -> Result<usize, OperationError> {
+        world.get_entity(self.source).or_broken()?
+            .buffered_count::<T>(session)
+    }
+
+    type Item = T;
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError> {
+        world.get_entity(self.source).or_broken()?
+            .clone_from_buffer(session)
+    }
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]> {
+        SmallVec::from_iter([self.source])
+    }
+}
+
+impl<T0, T1> Buffered for (T0, T1)
+where
+    T0: Buffered,
+    T1: Buffered,
+{
+    fn buffered_count(
+        &self,
+        session: Entity,
+        world: &World,
+    ) -> Result<usize, OperationError> {
+        Ok([
+            self.0.buffered_count(session, world)?,
+            self.1.buffered_count(session, world)?,
+        ].iter().copied().min().unwrap_or(0))
+    }
+
+    type Item = (T0::Item, T1::Item);
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError> {
+        let t0 = self.0.pull(session, world)?;
+        let t1 = self.1.pull(session, world)?;
+        Ok((t0, t1))
+    }
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]> {
+        let mut inputs = SmallVec::new();
+        inputs.extend(self.0.as_input());
+        inputs.extend(self.1.as_input());
+        inputs
+    }
+}
+
+impl<T0, T1, T2> Buffered for (T0, T1, T2)
+where
+    T0: Buffered,
+    T1: Buffered,
+    T2: Buffered,
+{
+    fn buffered_count(
+        &self,
+        session: Entity,
+        world: &World,
+    ) -> Result<usize, OperationError> {
+        Ok([
+            self.0.buffered_count(session, world)?,
+            self.1.buffered_count(session, world)?,
+            self.2.buffered_count(session, world)?,
+        ].iter().copied().min().unwrap_or(0))
+    }
+
+    type Item = (T0::Item, T1::Item, T2::Item);
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError> {
+        let t0 = self.0.pull(session, world)?;
+        let t1 = self.1.pull(session, world)?;
+        let t2 = self.2.pull(session, world)?;
+        Ok((t0, t1, t2))
+    }
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]> {
+        let mut inputs = SmallVec::new();
+        inputs.extend(self.0.as_input());
+        inputs.extend(self.1.as_input());
+        inputs.extend(self.2.as_input());
+        inputs
+    }
+}
+
+impl<T: Buffered, const N: usize> Buffered for [T; N] {
+    fn buffered_count(
+        &self,
+        session: Entity,
+        world: &World,
+    ) -> Result<usize, OperationError> {
+        let mut min_count = None;
+        for buffer in self.iter() {
+            let count = buffer.buffered_count(session, world)?;
+            if !min_count.is_some_and(|min| min < count) {
+                min_count = Some(count);
+            }
+        }
+
+        Ok(min_count.unwrap_or(0))
+    }
+
+    // TODO(@mxgrey) We may be able to use [T::Item; N] here instead of SmallVec
+    // when try_map is stabilized: https://github.com/rust-lang/rust/issues/79711
+    type Item = SmallVec<[T::Item; N]>;
+    fn pull(
+        &self,
+        session: Entity,
+        world: &mut World,
+    ) -> Result<Self::Item, OperationError> {
+        self.iter().map(|buffer| {
+            buffer.pull(session, world)
+        }).collect()
+    }
+
+    fn as_input(&self) -> SmallVec<[Entity; 8]> {
+        self.iter().flat_map(|buffer| buffer.as_input()).collect()
+    }
+}