Improve subtrace len compute algo (#186)

air/tests/test_module/issues/issue_173.rs

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2021 Fluence Labs Limited
+ *
+ * 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 air_test_utils::prelude::*;
+
+use fstrings::f;
+use fstrings::format_args_f;
+
+#[test]
+// test for github.com/fluencelabs/aquavm/issues/173
+fn issue_173() {
+    let set_variable_peer_id = "set_variable_peer_id";
+    let local_vm_peer_id_1 = "local_vm_peer_id_1";
+    let local_vm_peer_id_2 = "local_vm_peer_id_2";
+
+    let mut local_vm_1 = create_avm(echo_call_service(), local_vm_peer_id_1);
+    let mut local_vm_2 = create_avm(echo_call_service(), local_vm_peer_id_2);
+
+    let variables_mapping = maplit::hashmap! {
+        "1".to_string() => json!(1),
+        "2".to_string() => json!(2),
+    };
+    let mut set_variable_vm = create_avm(
+        set_variables_call_service(variables_mapping, VariableOptionSource::Argument(0)),
+        set_variable_peer_id,
+    );
+
+    let script = f!(r#"
+            (seq
+                (seq
+                    (call "{set_variable_peer_id}" ("" "") ["1"] $stream)
+                    (call "{set_variable_peer_id}" ("" "") ["2"] $stream)
+                )
+                (fold $stream i
+                    (par
+                        (new $stream
+                            (seq
+                                (seq
+                                    (call "{local_vm_peer_id_1}" ("" "") [i] $stream)
+                                    (next i)
+                                )
+                                (call "{local_vm_peer_id_1}" ("" "") [$stream])
+                            )
+                        )
+                        (call "{local_vm_peer_id_2}" ("" "") [$stream])
+                    )
+                )
+            )"#);
+
+    let result = checked_call_vm!(set_variable_vm, "", &script, "", "");
+    let vm_1_result = checked_call_vm!(local_vm_1, "", &script, "", result.data);
+    let vm_2_result = checked_call_vm!(local_vm_2, "", &script, "", vm_1_result.data.clone());
+
+    let vm_1_result = checked_call_vm!(local_vm_1, "", &script, vm_1_result.data, vm_2_result.data.clone());
+    let vm_2_result = checked_call_vm!(local_vm_2, "", script, vm_2_result.data, vm_1_result.data);
+
+    let actual_trace = trace_from_result(&vm_2_result);
+    let expected_trace = vec![
+        executed_state::stream_number(1, 0),
+        executed_state::stream_number(2, 0),
+        executed_state::fold(vec![
+            executed_state::subtrace_lore(0, SubTraceDesc::new(3, 2), SubTraceDesc::new(9, 2)),
+            executed_state::subtrace_lore(1, SubTraceDesc::new(5, 2), SubTraceDesc::new(7, 2)),
+        ]),
+        executed_state::par(6, 1),
+        executed_state::stream_number(1, 0),
+        executed_state::par(2, 1),
+        executed_state::stream_number(2, 0),
+        executed_state::scalar(json!([2])),
+        executed_state::scalar(json!([1, 2])),
+        executed_state::scalar(json!([1])),
+        executed_state::scalar(json!([1, 2])),
+    ];
+    assert_eq!(actual_trace, expected_trace);
+
+    let data = data_from_result(&vm_2_result);
+    let actual_restricted_streams = data.restricted_streams;
+    let expected_restricted_streams = maplit::hashmap! {
+        "$stream".to_string() => maplit::hashmap! {
+            282 => vec![1,1]
+        }
+    };
+    assert_eq!(actual_restricted_streams, expected_restricted_streams);
+}

air/tests/test_module/issues/mod.rs

@@ -15,6 +15,7 @@
  */
 
 mod issue_137;
+mod issue_173;
 mod issue_177;
 mod issue_178;
 mod issue_180;

crates/air-lib/trace-handler/src/data_keeper/trace_slider.rs

@@ -38,7 +38,7 @@ pub struct TraceSlider {
 }
 
 impl TraceSlider {
-    pub(super) fn new(trace: ExecutionTrace) -> Self {
+    pub(crate) fn new(trace: ExecutionTrace) -> Self {
         let subtrace_len = trace.len();
 
         Self {

crates/air-lib/trace-handler/src/merger/fold_merger/fold_lore_resolver.rs

@@ -64,12 +64,12 @@ pub(super) fn resolve_fold_lore(fold: &FoldResult, merge_ctx: &MergeCtx) -> Merg
 ///
 /// Imagine a fold on stream with 3 elements that have the same generation, in this case the
 /// conversion will look like this:
-/// [1, 1] [2, 2] [3, 3] => [6, 1] [5, 3] [3, 6]
+/// [1, 1] [2, 2] [3, 3] => [12, 1] [11, 3] [9, 6]
 ///   g0     g0     g0
 /// here a number before comma represents count of elements before next, and after the comma - after
 ///
 /// For fold with 5 elements of two generations:
-/// [1, 1] [2, 2] [3, 3] [4, 4] [5, 5] [1, 1] => [6, 1] [5, 3] [3, 6] [9, 4] [5, 9] [1, 1]
+/// [1, 1] [2, 2] [3, 3] [4, 4] [5, 5] [1, 1] => [12, 1] [11, 3] [9, 6] [18, 4] [14, 9] [2, 1]
 ///   g0     g0     g0     g1     g1     g2
 ///
 /// It could be seen that this function does a convolution of lens with respect to generations.
@@ -95,7 +95,7 @@ fn compute_lens_convolution(fold: &FoldResult, merge_ctx: &MergeCtx) -> MergeRes
         // TODO: check sequence for monotone
         if last_seen_generation != current_generation {
             if subtrace_id > 0 {
-                // do a back traversal for
+                // do a back traversal for before lens
                 compute_before_lens(&mut lens, last_seen_generation_pos, subtrace_id - 1);
             }
             last_seen_generation = current_generation;
@@ -130,13 +130,13 @@ fn compute_lens_convolution(fold: &FoldResult, merge_ctx: &MergeCtx) -> MergeRes
 
 fn compute_before_lens(lore_lens: &mut [LoresLen], begin_pos: usize, end_pos: usize) {
     let mut cum_before_len = 0;
+    let after_len = lore_lens[end_pos].after_len;
 
     for subtrace_id in (begin_pos..=end_pos).rev() {
-        let lens = &mut lore_lens[subtrace_id];
+        let before_len = &mut lore_lens[subtrace_id].before_len;
 
-        let current_before_len = lens.before_len;
+        cum_before_len += *before_len;
-        cum_before_len += current_before_len;
+        *before_len = cum_before_len + after_len;
-        lens.before_len = cum_before_len;
     }
 }
 
@@ -184,3 +184,127 @@ impl LoresLen {
         Self { before_len, after_len }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::compute_lens_convolution;
+    use crate::data_keeper::TraceSlider;
+    use crate::merger::fold_merger::fold_lore_resolver::LoresLen;
+    use crate::MergeCtx;
+    use air_interpreter_data::ApResult;
+    use air_interpreter_data::ExecutedState;
+    use air_interpreter_data::FoldResult;
+    use air_interpreter_data::FoldSubTraceLore;
+    use air_interpreter_data::SubTraceDesc;
+
+    #[test]
+    fn empty_fold_result() {
+        let lore = vec![];
+
+        let fold_result = FoldResult { lore };
+
+        let slider = TraceSlider::new(vec![]);
+        let ctx = MergeCtx {
+            slider,
+            streams: <_>::default(),
+        };
+
+        let (all_states, convoluted_lens) =
+            compute_lens_convolution(&fold_result, &ctx).expect("convolution should be successful");
+        assert_eq!(all_states, 0);
+        assert_eq!(convoluted_lens, vec![]);
+    }
+
+    #[test]
+    fn convolution_test_1() {
+        // [1, 1] [2, 2] [3, 3] => [12, 1] [11, 3] [9, 6]
+        //   g0     g0     g0
+        let lore = vec![
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 1), SubTraceDesc::new(0, 1)],
+            },
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 2), SubTraceDesc::new(0, 2)],
+            },
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 3), SubTraceDesc::new(0, 3)],
+            },
+        ];
+
+        let fold_result = FoldResult { lore };
+
+        let slider = TraceSlider::new(vec![ExecutedState::Ap(ApResult::new(vec![0]))]);
+        let ctx = MergeCtx {
+            slider,
+            streams: <_>::default(),
+        };
+
+        let (all_states, convoluted_lens) =
+            compute_lens_convolution(&fold_result, &ctx).expect("convolution should be successful");
+        assert_eq!(all_states, 12);
+
+        let expected_lens = vec![LoresLen::new(12, 1), LoresLen::new(11, 3), LoresLen::new(9, 6)];
+        assert_eq!(convoluted_lens, expected_lens);
+    }
+
+    #[test]
+    fn convolution_test_2() {
+        // [1, 1] [2, 2] [3, 3] [4, 4] [5, 5] [1, 1] => [12, 1] [11, 3] [9, 6] [18, 4] [14, 9] [2, 1]
+        //   g0     g0     g0     g1     g1     g2
+        let lore = vec![
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 1), SubTraceDesc::new(0, 1)],
+            },
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 2), SubTraceDesc::new(0, 2)],
+            },
+            FoldSubTraceLore {
+                value_pos: 0,
+                subtraces_desc: vec![SubTraceDesc::new(0, 3), SubTraceDesc::new(0, 3)],
+            },
+            FoldSubTraceLore {
+                value_pos: 1,
+                subtraces_desc: vec![SubTraceDesc::new(0, 4), SubTraceDesc::new(0, 4)],
+            },
+            FoldSubTraceLore {
+                value_pos: 1,
+                subtraces_desc: vec![SubTraceDesc::new(0, 5), SubTraceDesc::new(0, 5)],
+            },
+            FoldSubTraceLore {
+                value_pos: 2,
+                subtraces_desc: vec![SubTraceDesc::new(0, 1), SubTraceDesc::new(0, 1)],
+            },
+        ];
+
+        let fold_result = FoldResult { lore };
+
+        let slider = TraceSlider::new(vec![
+            ExecutedState::Ap(ApResult::new(vec![0])),
+            ExecutedState::Ap(ApResult::new(vec![1])),
+            ExecutedState::Ap(ApResult::new(vec![2])),
+        ]);
+        let ctx = MergeCtx {
+            slider,
+            streams: <_>::default(),
+        };
+
+        let (all_states, convoluted_lens) =
+            compute_lens_convolution(&fold_result, &ctx).expect("convolution should be successful");
+        assert_eq!(all_states, 32);
+
+        let expected_lens = vec![
+            LoresLen::new(12, 1),
+            LoresLen::new(11, 3),
+            LoresLen::new(9, 6),
+            LoresLen::new(18, 4),
+            LoresLen::new(14, 9),
+            LoresLen::new(2, 1),
+        ];
+        assert_eq!(convoluted_lens, expected_lens);
+    }
+}