Seems to work for tokio::spawn too

crates/paralegal-flow/src/ana/mod.rs

@@ -430,7 +430,10 @@ impl<'tcx> CallChangeCallback<'tcx> for MyCallback<'tcx> {
                 InlineJudgement::UseFlowModel(model) => {
                     // Set in case of errors
                     skip = SkipCall::Skip;
-                    assert!(matches!(model, FlowModel::SubClosure));
+                    assert!(matches!(
+                        model,
+                        FlowModel::SubClosure | FlowModel::SubFuture
+                    ));
                     if let [clj] = &info.arguments {
                         let ty = clj.ty(info.caller_body, self.tcx);
                         let (def_id, args) =
@@ -439,6 +442,15 @@ impl<'tcx> CallChangeCallback<'tcx> for MyCallback<'tcx> {
                             .unwrap()
                             .unwrap();
                         assert_eq!(instance.sig(self.tcx).unwrap().inputs().len(), 1);
+                        match model {
+                            FlowModel::SubClosure => {
+                                assert_eq!(
+                                    self.tcx.def_kind(def_id),
+                                    rustc_hir::def::DefKind::Closure
+                                )
+                            }
+                            FlowModel::SubFuture => assert!(self.tcx.generator_is_async(def_id)),
+                        };
                         skip = SkipCall::Replace {
                             instance,
                             calling_convention: CallingConvention::Indirect {

crates/paralegal-flow/src/args.rs

@@ -556,6 +556,8 @@ pub enum FlowModel {
     /// Replaces the result of a call to a higher-order function with a call to
     /// the input closure.
     SubClosure,
+    /// Replaces the result of a higher-order future by an input future.
+    SubFuture,
 }
 
 /// Additional configuration for the build process/rustc

crates/paralegal-flow/tests/flow-models.rs

@@ -30,3 +30,16 @@ define_test!(thread_spawn: graph -> {
     assert!(src.flows_to_data(&pass));
     assert!(pass.flows_to_data(&target));
 });
+
+define_test!(async_spawn: graph -> {
+    let src = graph.marked(Identifier::new_intern("source"));
+    let pass = graph.marked(Identifier::new_intern("pass"));
+    let target = graph.marked(Identifier::new_intern("target"));
+
+    assert!(!src.is_empty());
+    assert!(!pass.is_empty());
+    assert!(!target.is_empty());
+
+    assert!(src.flows_to_data(&pass));
+    assert!(pass.flows_to_data(&target));
+});

crates/paralegal-flow/tests/flow-models/Cargo.toml

@@ -8,3 +8,4 @@ edition = "2021"
 [dependencies]
 #async-std = "1"
 paralegal = { path = "../../../paralegal" }
+tokio = { version = "1", features = ["rt"] }

crates/paralegal-flow/tests/flow-models/Paralegal.toml

@@ -1,2 +1,6 @@
 [flow-models."std::thread::spawn"]
 mode = "sub-closure"
+
+
+[flow-models."tokio::spawn"]
+mode = "sub-future"

crates/paralegal-flow/tests/flow-models/src/main.rs

@@ -19,3 +19,10 @@ fn thread_spawn() {
 }
 
 fn main() {}
+
+#[paralegal::analyze]
+async fn async_spawn() {
+    let src = source();
+    let next = tokio::spawn(async move { pass(src) }).await.unwrap();
+    target(next);
+}