Update graph test

Author: akiradeveloper

csplib/tests/computational_graph.rs

@@ -1,40 +1,67 @@
 use csplib::*;
 
+#[csplib::process]
+struct P1 {
+    #[input]
+    a: i32,
+    #[output]
+    b: i32,
+}
+// λx. x+2
+async fn run_p1(inner: P1Inner) -> Result<()> {
+    let x = inner.a_r.reader().get().await?;
+    inner.b_w.put(x + 2)?;
+    Ok(())
+}
+#[csplib::process]
+struct P2 {
+    #[input]
+    a: i32,
+    #[output]
+    b: i32,
+}
+// λx. x*2
+async fn run_p2(inner: P2Inner) -> Result<()> {
+    let x = inner.a_r.reader().get().await?;
+    // Emulating expensive I/O
+    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+    inner.b_w.put(x * 2)?;
+    Ok(())
+}
+#[csplib::process]
+struct P3 {
+    #[input]
+    a: i32,
+    #[input]
+    b: i32,
+    #[output]
+    c: i32,
+}
+// λxy. x*y
+async fn run_p3(inner: P3Inner) -> Result<()> {
+    let x = inner.a_r.reader().get().await?;
+    let y = inner.b_r.reader().get().await?;
+    inner.c_w.put(x * y)?;
+    Ok(())
+}
+
 #[tokio::test]
 async fn computational_graph() {
-    let (w1, ch1) = channel();
-    let (w2, ch2) = channel();
-    let (w3, ch3) = channel();
-    let (w4, ch4) = channel();
-    // λx. x+2
-    tokio::spawn({
-        let r1 = ch1.reader();
-        async move {
-            let x = r1.get().await.unwrap();
-            w2.put(x + 2).unwrap();
-        }
-    });
-    // λx. x*2
-    tokio::spawn({
-        let r1 = ch1.reader();
-        async move {
-            // Emulating expensive I/O
-            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
-            let x = r1.get().await.unwrap();
-            w3.put(x * 2).unwrap();
-        }
-    });
-    // λxy. x*y
-    tokio::spawn({
-        let r2 = ch2.reader();
-        let r3 = ch3.reader();
-        async move {
-            let (x, y) = tokio::try_join!(r2.get(), r3.get()).unwrap();
-            w4.put(x * y).unwrap();
-        }
-    });
-    w1.put(1).unwrap();
-    let r4 = ch4.reader();
-    let ans = r4.get().await.unwrap();
+    let (main_w, main_r) = channel();
+    let (p1, p1_inner) = P1::new();
+    let (p2, p2_inner) = P2::new();
+    let (p3, p3_inner) = P3::new();
+
+    tokio::spawn(run_p1(p1_inner));
+    tokio::spawn(run_p2(p2_inner));
+    tokio::spawn(run_p3(p3_inner));
+
+    tokio::spawn(connect(main_r.reader(), p1.a_w));
+    tokio::spawn(connect(main_r.reader(), p2.a_w));
+    tokio::spawn(connect(p1.b_r.reader(), p3.a_w));
+    tokio::spawn(connect(p2.b_r.reader(), p3.b_w));
+
+    main_w.put(1).unwrap();
+    let ans = p3.c_r.reader().get().await.unwrap();
     assert_eq!(ans, 6);
 }