stablize the order of outputs passed to join_vec

Signed-off-by: Teo Koon Peng <[email protected]>

Author: koonpeng

src/diagram.rs

@@ -40,14 +40,14 @@ pub type OperationId = String;
 )]
 #[serde(untagged, rename_all = "snake_case")]
 pub enum NextOperation {
-    Target(String),
+    Target(OperationId),
     Builtin { builtin: BuiltinTarget },
 }
 
 impl Display for NextOperation {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::Target(builder_id) => f.write_str(builder_id),
+            Self::Target(operation_id) => f.write_str(operation_id),
             Self::Builtin { builtin } => write!(f, "builtin:{}", builtin),
         }
     }
@@ -77,6 +77,49 @@ pub enum BuiltinTarget {
     Dispose,
 }
 
+#[derive(
+    Debug, Clone, Serialize, Deserialize, JsonSchema, Hash, PartialEq, Eq, PartialOrd, Ord,
+)]
+#[serde(untagged, rename_all = "snake_case")]
+pub enum SourceOperation {
+    Source(OperationId),
+    Builtin { builtin: BuiltinSource },
+}
+
+impl From<OperationId> for SourceOperation {
+    fn from(value: OperationId) -> Self {
+        SourceOperation::Source(value)
+    }
+}
+
+impl Display for SourceOperation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Source(operation_id) => f.write_str(operation_id),
+            Self::Builtin { builtin } => write!(f, "builtin:{}", builtin),
+        }
+    }
+}
+
+#[derive(
+    Debug,
+    Clone,
+    Serialize,
+    Deserialize,
+    JsonSchema,
+    Hash,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    strum::Display,
+)]
+#[serde(rename_all = "snake_case")]
+#[strum(serialize_all = "snake_case")]
+pub enum BuiltinSource {
+    Start,
+}
+
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub struct TerminateOp {}

src/diagram/join.rs

@@ -7,13 +7,19 @@ use tracing::debug;
 
 use crate::{Builder, IterBufferable, Output};
 
-use super::{DiagramError, DynOutput, NextOperation, NodeRegistry, SerializeMessage};
+use super::{
+    DiagramError, DynOutput, NextOperation, NodeRegistry, SerializeMessage, SourceOperation,
+};
 
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub struct JoinOp {
     pub(super) next: NextOperation,
 
+    /// Controls the order of the resulting join. Each item must be an operation id of one of the
+    /// incoming outputs.
+    pub(super) order: Vec<SourceOperation>,
+
     /// Do not serialize before performing the join. If true, joins can only be done
     /// on outputs of the same type.
     pub(super) no_serialize: Option<bool>,
@@ -31,11 +37,6 @@ where
     registry
         .join_impls
         .insert(TypeId::of::<T>(), Box::new(join_impl::<T>));
-
-    // FIXME(koonpeng): join_vec results in a SmallVec<[T; N]>, we can't serialize it because
-    // it doesn't implement JsonSchema, and we can't impl it because of orphan rule. We would need
-    // to create our own trait that covers `JsonSchema` and `Serialize`.
-    // register_serialize::<Vec<T>, Serializer>(registry);
 }
 
 /// Serialize the outputs before joining them, and convert the resulting joined output into a
@@ -151,9 +152,9 @@ mod tests {
             "ops": {
                 "split": {
                     "type": "split",
-                    "index": ["getSplitValue1", "getSplitValue2"]
+                    "index": ["get_split_value1", "get_split_value2"]
                 },
-                "getSplitValue1": {
+                "get_split_value1": {
                     "type": "node",
                     "builder": "get_split_value",
                     "next": "op1",
@@ -163,7 +164,7 @@ mod tests {
                     "builder": "multiply3_uncloneable",
                     "next": "join",
                 },
-                "getSplitValue2": {
+                "get_split_value2": {
                     "type": "node",
                     "builder": "get_split_value",
                     "next": "op2",
@@ -175,10 +176,11 @@ mod tests {
                 },
                 "join": {
                     "type": "join",
-                    "next": "serializeJoinOutput",
+                    "order": ["op1", "op2"],
+                    "next": "serialize_join_output",
                     "no_serialize": true,
                 },
-                "serializeJoinOutput": {
+                "serialize_join_output": {
                     "type": "node",
                     "builder": "serialize_join_output",
                     "next": { "builtin": "terminate" },
@@ -191,10 +193,8 @@ mod tests {
             .spawn_and_run(&diagram, serde_json::Value::from([1, 2]))
             .unwrap();
         assert_eq!(result.as_array().unwrap().len(), 2);
-        // order is not guaranteed so need to test for both possibility
-        assert!(result[0] == 3 || result[0] == 6);
-        assert!(result[1] == 3 || result[1] == 6);
-        assert!(result[0] != result[1]);
+        assert_eq!(result[0], 3);
+        assert_eq!(result[1], 6);
     }
 
     #[test]
@@ -241,6 +241,7 @@ mod tests {
                 },
                 "join": {
                     "type": "join",
+                    "order": [],
                     "next": { "builtin": "terminate" },
                     "no_serialize": true,
                 },
@@ -296,6 +297,7 @@ mod tests {
                 },
                 "join": {
                     "type": "join",
+                    "order": ["op1", "op2"],
                     "next": { "builtin": "terminate" },
                 },
             }
@@ -306,9 +308,7 @@ mod tests {
             .spawn_and_run(&diagram, serde_json::Value::Null)
             .unwrap();
         assert_eq!(result.as_array().unwrap().len(), 2);
-        // order is not guaranteed so need to test for both possibility
-        assert!(result[0] == 1 || result[0] == "hello");
-        assert!(result[1] == 1 || result[1] == "hello");
-        assert!(result[0] != result[1]);
+        assert_eq!(result[0], 1);
+        assert_eq!(result[1], "hello");
     }
 }

src/diagram/workflow_builder.rs

@@ -7,7 +7,7 @@ use crate::{diagram::join::serialize_and_join, Builder, InputSlot, Output, Strea
 use super::{
     fork_clone::DynForkClone, impls::DefaultImpl, split_chain, transform::transform_output,
     BuiltinTarget, Diagram, DiagramError, DiagramOperation, DiagramScope, DynInputSlot, DynOutput,
-    NextOperation, NodeOp, NodeRegistry, OperationId, SplitOpParams,
+    NextOperation, NodeOp, NodeRegistry, OperationId, SourceOperation, SplitOpParams,
 };
 
 struct Vertex<'a> {
@@ -18,8 +18,7 @@ struct Vertex<'a> {
 }
 
 struct Edge<'a> {
-    /// The source of the edge, may be `None` if it comes from outside the diagram, e.g. the entry point of the diagram.
-    source: Option<&'a OperationId>,
+    source: SourceOperation,
     target: &'a NextOperation,
     state: EdgeState<'a>,
 }
@@ -79,7 +78,9 @@ pub(super) fn create_workflow<'a, Streams: StreamPack>(
             edges.insert(
                 edges.len(),
                 Edge {
-                    source: None,
+                    source: SourceOperation::Builtin {
+                        builtin: super::BuiltinSource::Start,
+                    },
                     target: &diagram.start,
                     state: EdgeState::Ready {
                         output: scope.input.into(),
@@ -95,10 +96,16 @@ pub(super) fn create_workflow<'a, Streams: StreamPack>(
 
     let mut terminate_edges: Vec<usize> = Vec::new();
 
-    let mut add_edge = |source: Option<&'a OperationId>,
+    let mut add_edge = |source: SourceOperation,
                         target: &'a NextOperation,
                         state: EdgeState<'a>|
      -> Result<(), DiagramError> {
+        let source_id = if let SourceOperation::Source(source) = &source {
+            Some(source.clone())
+        } else {
+            None
+        };
+
         edges.insert(
             edges.len(),
             Edge {
@@ -109,10 +116,10 @@ pub(super) fn create_workflow<'a, Streams: StreamPack>(
         );
         let new_edge_id = edges.len() - 1;
 
-        if let Some(source) = source {
+        if let Some(source_id) = source_id {
             let source_vertex = vertices
-                .get_mut(source)
-                .ok_or_else(|| DiagramError::OperationNotFound(source.clone()))?;
+                .get_mut(&source_id)
+                .ok_or_else(|| DiagramError::OperationNotFound(source_id.clone()))?;
             source_vertex.out_edges.push(new_edge_id);
         }
 
@@ -141,7 +148,7 @@ pub(super) fn create_workflow<'a, Streams: StreamPack>(
                 let n = reg.create_node(builder, node_op.config.clone())?;
                 inputs.insert(op_id, n.input);
                 add_edge(
-                    Some(op_id),
+                    op_id.clone().into(),
                     &node_op.next,
                     EdgeState::Ready {
                         output: n.output.into(),
@@ -151,35 +158,39 @@ pub(super) fn create_workflow<'a, Streams: StreamPack>(
             }
             DiagramOperation::ForkClone(fork_clone_op) => {
                 for next_op_id in fork_clone_op.next.iter() {
-                    add_edge(Some(op_id), next_op_id, EdgeState::Pending)?;
+                    add_edge(op_id.clone().into(), next_op_id, EdgeState::Pending)?;
                 }
             }
             DiagramOperation::Unzip(unzip_op) => {
                 for next_op_id in unzip_op.next.iter() {
-                    add_edge(Some(op_id), next_op_id, EdgeState::Pending)?;
+                    add_edge(op_id.clone().into(), next_op_id, EdgeState::Pending)?;
                 }
             }
             DiagramOperation::ForkResult(fork_result_op) => {
-                add_edge(Some(op_id), &fork_result_op.ok, EdgeState::Pending)?;
-                add_edge(Some(op_id), &fork_result_op.err, EdgeState::Pending)?;
+                add_edge(op_id.clone().into(), &fork_result_op.ok, EdgeState::Pending)?;
+                add_edge(
+                    op_id.clone().into(),
+                    &fork_result_op.err,
+                    EdgeState::Pending,
+                )?;
             }
             DiagramOperation::Split(split_op) => {
                 let next_op_ids: Vec<&NextOperation> = match &split_op.params {
                     SplitOpParams::Index(v) => v.iter().collect(),
                     SplitOpParams::Key(v) => v.values().collect(),
                 };
                 for next_op_id in next_op_ids {
-                    add_edge(Some(op_id), next_op_id, EdgeState::Pending)?;
+                    add_edge(op_id.clone().into(), next_op_id, EdgeState::Pending)?;
                 }
                 if let Some(remaining) = &split_op.remaining {
-                    add_edge(Some(op_id), &remaining, EdgeState::Pending)?;
+                    add_edge(op_id.clone().into(), &remaining, EdgeState::Pending)?;
                 }
             }
             DiagramOperation::Join(join_op) => {
-                add_edge(Some(op_id), &join_op.next, EdgeState::Pending)?;
+                add_edge(op_id.clone().into(), &join_op.next, EdgeState::Pending)?;
             }
             DiagramOperation::Transform(transform_op) => {
-                add_edge(Some(op_id), &transform_op.next, EdgeState::Pending)?;
+                add_edge(op_id.clone().into(), &transform_op.next, EdgeState::Pending)?;
             }
             DiagramOperation::Dispose => {}
         }
@@ -252,23 +263,31 @@ fn connect_vertex<'a>(
             if target.in_edges.is_empty() {
                 return Err(DiagramError::EmptyJoin);
             }
-            let outputs: Vec<DynOutput> = target
+            let mut outputs: HashMap<SourceOperation, DynOutput> = target
                 .in_edges
                 .iter()
                 .map(|e| {
                     let edge = edges.remove(e).unwrap();
                     match edge.state {
-                        EdgeState::Ready { output, origin: _ } => output,
+                        EdgeState::Ready { output, origin: _ } => (edge.source, output),
                         _ => panic!("expected all incoming edges to be ready"),
                     }
                 })
                 .collect();
 
+            let mut ordered_outputs: Vec<DynOutput> = Vec::with_capacity(target.in_edges.len());
+            for source_id in join_op.order.iter() {
+                let o = outputs
+                    .remove(source_id)
+                    .ok_or(DiagramError::OperationNotFound(source_id.to_string()))?;
+                ordered_outputs.push(o);
+            }
+
             let joined_output = if join_op.no_serialize.unwrap_or(false) {
-                let join_impl = &registry.join_impls[&outputs[0].type_id];
-                join_impl(builder, outputs)?
+                let join_impl = &registry.join_impls[&ordered_outputs[0].type_id];
+                join_impl(builder, ordered_outputs)?
             } else {
-                serialize_and_join(builder, registry, outputs)?.into()
+                serialize_and_join(builder, registry, ordered_outputs)?.into()
             };
 
             let out_edge = edges.get_mut(&target.out_edges[0]).unwrap();