-
Notifications
You must be signed in to change notification settings - Fork 37
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(dfir_lang): add
repeat_n windowing operator
- Loading branch information
1 parent
5e8e2f1
commit 59fa7e0
Showing
14 changed files
with
529 additions
and
32 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,57 @@ | ||
| use quote::quote_spanned; | ||
|
|
||
| use super::{OperatorConstraints, OperatorWriteOutput, WriteContextArgs}; | ||
|
|
||
| /// TODO(mingwei): docs | ||
| pub const REPEAT_N: OperatorConstraints = OperatorConstraints { | ||
| name: "repeat_n", | ||
| num_args: 1, | ||
| write_fn: |wc @ &WriteContextArgs { | ||
| context, | ||
| hydroflow, | ||
| op_span, | ||
| arguments, | ||
| .. | ||
| }, | ||
| diagnostics| { | ||
| let OperatorWriteOutput { | ||
| write_prologue, | ||
| write_iterator, | ||
| write_iterator_after, | ||
| } = (super::all_once::ALL_ONCE.write_fn)(wc, diagnostics)?; | ||
|
|
||
| let count_ident = wc.make_ident("count"); | ||
|
|
||
| let write_prologue = quote_spanned! {op_span=> | ||
| #write_prologue | ||
|
|
||
| let #count_ident = #hydroflow.add_state(::std::cell::Cell::new(0_usize)); | ||
| #hydroflow.set_state_tick_hook(#count_ident, move |cell| { cell.take(); }); | ||
| }; | ||
|
|
||
| // Reschedule, to repeat. | ||
| let count_arg = &arguments[0]; | ||
| let write_iterator_after = quote_spanned! {op_span=> | ||
| #write_iterator_after | ||
|
|
||
| { | ||
| let count_ref = #context.state_ref(#count_ident); | ||
| if #context.is_first_loop_iteration() { | ||
| count_ref.set(0); | ||
| } | ||
| let count = count_ref.get() + 1; | ||
| if count < #count_arg { | ||
| count_ref.set(count); | ||
| #context.reschedule_current_subgraph(); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| Ok(OperatorWriteOutput { | ||
| write_prologue, | ||
| write_iterator, | ||
| write_iterator_after, | ||
| }) | ||
| }, | ||
| ..super::all_once::ALL_ONCE | ||
| }; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
78 changes: 78 additions & 0 deletions
78
dfir_rs/tests/snapshots/surface_loop__flo_repeat_n@graphvis_dot.snap
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| --- | ||
| source: dfir_rs/tests/surface_loop.rs | ||
| expression: "df.meta_graph().unwrap().to_dot(& Default :: default())" | ||
| --- | ||
| digraph { | ||
| node [fontname="Monaco,Menlo,Consolas,"Droid Sans Mono",Inconsolata,"Courier New",monospace", style=filled]; | ||
| edge [fontname="Monaco,Menlo,Consolas,"Droid Sans Mono",Inconsolata,"Courier New",monospace"]; | ||
| n1v1 [label="(n1v1) source_iter([\"alice\", \"bob\"])", shape=invhouse, fillcolor="#88aaff"] | ||
| n2v1 [label="(n2v1) source_stream(iter_batches_stream(0..12, 3))", shape=invhouse, fillcolor="#88aaff"] | ||
| n3v1 [label="(n3v1) batch()", shape=invhouse, fillcolor="#88aaff"] | ||
| n4v1 [label="(n4v1) flatten()", shape=invhouse, fillcolor="#88aaff"] | ||
| n5v1 [label="(n5v1) batch()", shape=invhouse, fillcolor="#88aaff"] | ||
| n6v1 [label="(n6v1) flatten()", shape=invhouse, fillcolor="#88aaff"] | ||
| n7v1 [label="(n7v1) cross_join::<'static, 'tick>()", shape=invhouse, fillcolor="#88aaff"] | ||
| n8v1 [label="(n8v1) repeat_n(3)", shape=invhouse, fillcolor="#88aaff"] | ||
| n9v1 [label="(n9v1) map(|vec| (context.current_tick().0, vec))", shape=invhouse, fillcolor="#88aaff"] | ||
| n10v1 [label="(n10v1) inspect(|x| println!(\"{:?}\", x))", shape=invhouse, fillcolor="#88aaff"] | ||
| n11v1 [label="(n11v1) assert_eq([\l (\l 0,\l vec![\l (\"alice\", 0),\l (\"alice\", 1),\l (\"alice\", 2),\l (\"bob\", 0),\l (\"bob\", 1),\l (\"bob\", 2),\l ],\l ),\l (\l 0,\l vec![\l (\"alice\", 0),\l (\"alice\", 1),\l (\"alice\", 2),\l (\"bob\", 0),\l (\"bob\", 1),\l (\"bob\", 2),\l ],\l ),\l (\l 0,\l vec![\l (\"alice\", 0),\l (\"alice\", 1),\l (\"alice\", 2),\l (\"bob\", 0),\l (\"bob\", 1),\l (\"bob\", 2),\l ],\l ),\l (\l 1,\l vec![\l (\"alice\", 3),\l (\"alice\", 4),\l (\"alice\", 5),\l (\"bob\", 3),\l (\"bob\", 4),\l (\"bob\", 5),\l ],\l ),\l (\l 1,\l vec![\l (\"alice\", 3),\l (\"alice\", 4),\l (\"alice\", 5),\l (\"bob\", 3),\l (\"bob\", 4),\l (\"bob\", 5),\l ],\l ),\l (\l 1,\l vec![\l (\"alice\", 3),\l (\"alice\", 4),\l (\"alice\", 5),\l (\"bob\", 3),\l (\"bob\", 4),\l (\"bob\", 5),\l ],\l ),\l (\l 2,\l vec![\l (\"alice\", 6),\l (\"alice\", 7),\l (\"alice\", 8),\l (\"bob\", 6),\l (\"bob\", 7),\l (\"bob\", 8),\l ],\l ),\l (\l 2,\l vec![\l (\"alice\", 6),\l (\"alice\", 7),\l (\"alice\", 8),\l (\"bob\", 6),\l (\"bob\", 7),\l (\"bob\", 8),\l ],\l ),\l (\l 2,\l vec![\l (\"alice\", 6),\l (\"alice\", 7),\l (\"alice\", 8),\l (\"bob\", 6),\l (\"bob\", 7),\l (\"bob\", 8),\l ],\l ),\l (\l 3,\l vec![\l (\"alice\", 9),\l (\"alice\", 10),\l (\"alice\", 11),\l (\"bob\", 9),\l (\"bob\", 10),\l (\"bob\", 11),\l ],\l ),\l (\l 3,\l vec![\l (\"alice\", 9),\l (\"alice\", 10),\l (\"alice\", 11),\l (\"bob\", 9),\l (\"bob\", 10),\l (\"bob\", 11),\l ],\l ),\l (\l 3,\l vec![\l (\"alice\", 9),\l (\"alice\", 10),\l (\"alice\", 11),\l (\"bob\", 9),\l (\"bob\", 10),\l (\"bob\", 11),\l ],\l ),\l])\l", shape=house, fillcolor="#ffff88"] | ||
| n12v1 [label="(n12v1) handoff", shape=parallelogram, fillcolor="#ddddff"] | ||
| n13v1 [label="(n13v1) handoff", shape=parallelogram, fillcolor="#ddddff"] | ||
| n14v1 [label="(n14v1) handoff", shape=parallelogram, fillcolor="#ddddff"] | ||
| n4v1 -> n7v1 [label="0"] | ||
| n3v1 -> n4v1 | ||
| n1v1 -> n12v1 | ||
| n6v1 -> n7v1 [label="1"] | ||
| n5v1 -> n6v1 | ||
| n2v1 -> n13v1 | ||
| n10v1 -> n11v1 | ||
| n9v1 -> n10v1 | ||
| n8v1 -> n9v1 | ||
| n7v1 -> n14v1 | ||
| n12v1 -> n3v1 | ||
| n13v1 -> n5v1 | ||
| n14v1 -> n8v1 [color=red] | ||
| subgraph "cluster n1v1" { | ||
| fillcolor="#dddddd" | ||
| style=filled | ||
| label = "sg_1v1\nstratum 0" | ||
| n1v1 | ||
| subgraph "cluster_sg_1v1_var_users" { | ||
| label="var users" | ||
| n1v1 | ||
| } | ||
| } | ||
| subgraph "cluster n2v1" { | ||
| fillcolor="#dddddd" | ||
| style=filled | ||
| label = "sg_2v1\nstratum 0" | ||
| n2v1 | ||
| subgraph "cluster_sg_2v1_var_messages" { | ||
| label="var messages" | ||
| n2v1 | ||
| } | ||
| } | ||
| subgraph "cluster n3v1" { | ||
| fillcolor="#dddddd" | ||
| style=filled | ||
| label = "sg_3v1\nstratum 0" | ||
| n3v1 | ||
| n4v1 | ||
| n5v1 | ||
| n6v1 | ||
| n7v1 | ||
| subgraph "cluster_sg_3v1_var_cp" { | ||
| label="var cp" | ||
| n7v1 | ||
| } | ||
| } | ||
| subgraph "cluster n4v1" { | ||
| fillcolor="#dddddd" | ||
| style=filled | ||
| label = "sg_4v1\nstratum 1" | ||
| n8v1 | ||
| n9v1 | ||
| n10v1 | ||
| n11v1 | ||
| } | ||
| } |
Oops, something went wrong.