Fix ap trace handler behaviour (#224)

air/src/execution_step/air/ap.rs

@@ -23,6 +23,7 @@ use super::TraceHandler;
 use crate::execution_step::air::ValueAggregate;
 use crate::execution_step::boxed_value::Variable;
 use crate::execution_step::resolver::apply_lambda;
+use crate::log_instruction;
 use crate::trace_to_exec_err;
 use crate::JValue;
 use crate::SecurityTetraplet;
@@ -37,6 +38,7 @@ use std::rc::Rc;
 
 impl<'i> super::ExecutableInstruction<'i> for Ap<'i> {
     fn execute(&self, exec_ctx: &mut ExecutionCtx<'i>, trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
+        log_instruction!(call, exec_ctx, trace_ctx);
         let should_touch_trace = should_touch_trace(self);
         // this applying should be at the very beginning of this function,
         // because it's necessary to check argument lambda, for more details see

air/tests/test_module/issues/issue_221.rs

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2022 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/221
+fn issue_221() {
+    let peer_1_id = "peer_1_id";
+    let peer_2_id = "peer_2_id";
+    let join_1_id = "join_1_id";
+    let join_2_id = "join_2_id";
+    let set_variable_id = "set_variable_id";
+
+    let peer_1_value = "peer_1_value";
+    let peer_2_value = "peer_2_value";
+    let mut peer_1 = create_avm(set_variable_call_service(json!(peer_1_value)), peer_1_id);
+    let mut peer_2 = create_avm(set_variable_call_service(json!(peer_2_value)), peer_2_id);
+    let mut join_1 = create_avm(echo_call_service(), join_1_id);
+    let mut set_variable = create_avm(
+        set_variable_call_service(json!([peer_1_id, peer_2_id])),
+        set_variable_id,
+    );
+
+    let script = f!(r#"
+        (seq
+            (seq
+                (seq
+                    ;; let's peers be an array of two values [peer_1_id, peer_2_id]
+                    (call "{set_variable_id}" ("" "") [] peers)
+                    (fold peers peer
+                        (par
+                            (seq
+                                (call peer ("" "") [] value)
+                                ;; it's crucial to reproduce this bug to add value to stream
+                                ;; with help of ap instruction
+                                (ap value $stream)
+                            )
+                            (next peer)
+                        )
+                    )
+                )
+                ;; join streams on join_1/join_2 peers in such a way that will have different state:
+                ;; join_1 $stream: [peer_1_value, peer_2_value]
+                ;; join_2 $stream: [peer_2_value, peer_1_value]
+                (fold $stream iterator
+                    ;; here it'll encounter a bug in trace handler, because fold won't shuffle lores in
+                    ;; appropriate way and state for (1) is returned
+                    (par
+                        (par
+                            (call "{join_1_id}" ("" "") [iterator])
+                            (call "{join_2_id}" ("" "") [iterator])
+                        )
+                        (next iterator)
+                    )
+                )
+            )
+            (call "some_peer_id" ("" "") []) ;; (1)
+        )
+    "#);
+
+    let result = checked_call_vm!(set_variable, "", &script, "", "");
+    let peer_1_result = checked_call_vm!(peer_1, "", &script, "", result.data.clone());
+    let peer_2_result = checked_call_vm!(peer_2, "", &script, "", result.data.clone());
+
+    let join_1_result = checked_call_vm!(join_1, "", &script, "", peer_1_result.data.clone());
+    let join_1_result = checked_call_vm!(join_1, "", &script, join_1_result.data, peer_2_result.data.clone()); // before 0.20.9 it fails here
+    let actual_trace = trace_from_result(&join_1_result);
+    let expected_trace = vec![
+        executed_state::scalar(json!([peer_1_id, peer_2_id])),
+        executed_state::par(2, 3),
+        executed_state::scalar_string(peer_1_value),
+        executed_state::ap(Some(0)),
+        executed_state::par(2, 0),
+        executed_state::scalar_string(peer_2_value),
+        executed_state::ap(Some(1)),
+        executed_state::fold(vec![
+            executed_state::subtrace_lore(3, SubTraceDesc::new(8, 4), SubTraceDesc::new(12, 0)),
+            executed_state::subtrace_lore(6, SubTraceDesc::new(12, 4), SubTraceDesc::new(16, 0)),
+        ]),
+        executed_state::par(3, 0),
+        executed_state::par(1, 1),
+        executed_state::scalar_string(peer_1_value),
+        executed_state::request_sent_by(peer_1_id),
+        executed_state::par(3, 0),
+        executed_state::par(1, 1),
+        executed_state::scalar_string(peer_2_value),
+        executed_state::request_sent_by(peer_2_id),
+        executed_state::request_sent_by(join_1_id),
+    ];
+
+    assert_eq!(actual_trace, expected_trace);
+}

air/tests/test_module/issues/mod.rs

@@ -23,3 +23,4 @@ mod issue_206;
 mod issue_211;
 mod issue_214;
 mod issue_216;
+mod issue_221;

crates/air-lib/trace-handler/src/merger/ap_merger.rs

@@ -16,6 +16,8 @@
 
 use super::*;
 
+const EXPECTED_STATE_NAME: &str = "ap";
+
 #[derive(Debug, Clone)]
 pub enum MergerApResult {
     /// There is no corresponding state in a trace for this call.
@@ -28,20 +30,37 @@ pub enum MergerApResult {
 
 pub(crate) fn try_merge_next_state_as_ap(data_keeper: &mut DataKeeper) -> MergeResult<MergerApResult> {
     use ExecutedState::Ap;
+    use PreparationScheme::*;
 
     let prev_state = data_keeper.prev_slider_mut().next_state();
     let current_state = data_keeper.current_slider_mut().next_state();
 
-    let ap = match (prev_state, current_state) {
-        (Some(Ap(prev_ap)), _) => prev_ap,
+    match (prev_state, current_state) {
+        (Some(Ap(prev_ap)), Some(_)) => prepare_merge_result(Some(prev_ap), Both, data_keeper),
+        (Some(Ap(prev_ap)), None) => prepare_merge_result(Some(prev_ap), Previous, data_keeper),
         // check that current state is Ap, but it's impossible to use it, because prev_data
         // could not have streams with such generations
-        (None, Some(Ap(_))) => return Ok(MergerApResult::Empty),
-        (None, None) => return Ok(MergerApResult::Empty),
-        (prev_state, current_state) => return Err(MergeError::incompatible_states(prev_state, current_state, "ap")),
-    };
+        (None, Some(Ap(_))) => prepare_merge_result(None, Current, data_keeper),
+        (None, None) => Ok(MergerApResult::Empty),
+        (prev_state, current_state) => Err(MergeError::incompatible_states(
+            prev_state,
+            current_state,
+            EXPECTED_STATE_NAME,
+        )),
+    }
+}
 
-    to_merger_result(ap)
+fn prepare_merge_result(
+    ap_result: Option<ApResult>,
+    scheme: PreparationScheme,
+    data_keeper: &mut DataKeeper,
+) -> MergeResult<MergerApResult> {
+    prepare_positions_mapping(scheme, data_keeper);
+
+    match ap_result {
+        Some(ap_result) => to_merger_result(ap_result),
+        None => Ok(MergerApResult::Empty),
+    }
 }
 
 macro_rules! to_maybe_generation {

crates/air-lib/trace-handler/src/merger/call_merger.rs

@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-mod call_result_constructor;
 mod utils;
 
 use super::*;
 use air_parser::ast::CallOutputValue;
-use call_result_constructor::*;
 use utils::*;
 
+const EXPECTED_STATE_NAME: &str = "call";
+
 #[derive(Debug, Clone)]
 pub enum MergerCallResult {
     /// There is no corresponding state in a trace for this call.
@@ -37,7 +37,7 @@ pub(crate) fn try_merge_next_state_as_call(
     output_value: &CallOutputValue<'_>,
 ) -> MergeResult<MergerCallResult> {
     use ExecutedState::Call;
-    use PrepareScheme::*;
+    use PreparationScheme::*;
 
     let prev_state = data_keeper.prev_slider_mut().next_state();
     let current_state = data_keeper.current_slider_mut().next_state();
@@ -53,7 +53,13 @@ pub(crate) fn try_merge_next_state_as_call(
         (None, Some(Call(current_call))) => return Ok(prepare_call_result(current_call, Current, data_keeper)),
         (Some(Call(prev_call)), None) => return Ok(prepare_call_result(prev_call, Previous, data_keeper)),
         (None, None) => return Ok(MergerCallResult::Empty),
-        (prev_state, current_state) => return Err(MergeError::incompatible_states(prev_state, current_state, "call")),
+        (prev_state, current_state) => {
+            return Err(MergeError::incompatible_states(
+                prev_state,
+                current_state,
+                EXPECTED_STATE_NAME,
+            ))
+        }
     };
 
     let merged_call = merge_call_result(prev_call, current_call, value_type, data_keeper)?;
@@ -91,6 +97,17 @@ fn merge_call_result(
     Ok(merged_state)
 }
 
+pub(super) fn prepare_call_result(
+    value: CallResult,
+    scheme: PreparationScheme,
+    data_keeper: &mut DataKeeper,
+) -> MergerCallResult {
+    let trace_pos = data_keeper.result_states_count();
+    prepare_positions_mapping(scheme, data_keeper);
+
+    MergerCallResult::CallResult { value, trace_pos }
+}
+
 #[derive(Debug, Copy, Clone)]
 pub(crate) enum ValueType<'i> {
     Scalar,

crates/air-lib/trace-handler/src/merger/mod.rs

@@ -19,6 +19,7 @@ mod call_merger;
 mod errors;
 mod fold_merger;
 mod par_merger;
+mod position_mapping;
 
 pub use ap_merger::MergerApResult;
 pub use call_merger::MergerCallResult;
@@ -38,6 +39,9 @@ pub(super) use call_merger::try_merge_next_state_as_call;
 pub(crate) use fold_merger::try_merge_next_state_as_fold;
 pub(crate) use par_merger::try_merge_next_state_as_par;
 
+use position_mapping::prepare_positions_mapping;
+use position_mapping::PreparationScheme;
+
 type MergeResult<T> = std::result::Result<T, MergeError>;
 
 use super::data_keeper::DataPositions;

crates/air-lib/trace-handler/src/merger/position_mapping.rs

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 Fluence Labs Limited
+ * Copyright 2022 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.
@@ -14,33 +14,32 @@
  * limitations under the License.
  */
 
-use super::*;
+use super::DataKeeper;
+use super::DataPositions;
 
-pub(super) enum PrepareScheme {
+pub(super) enum PreparationScheme {
     Previous,
     Current,
     Both,
 }
 
-pub(super) fn prepare_call_result(
-    value: CallResult,
-    scheme: PrepareScheme,
-    data_keeper: &mut DataKeeper,
-) -> MergerCallResult {
+/// Prepares new_to_old_pos mapping in data keeper to keep track of value sources.
+pub(super) fn prepare_positions_mapping(scheme: PreparationScheme, data_keeper: &mut DataKeeper) {
+    use PreparationScheme::*;
+
+    // it's safe to sub 1 from positions iff scheme was set correctly
     let prev_pos = match scheme {
-        PrepareScheme::Previous | PrepareScheme::Both => Some(data_keeper.prev_slider().position() - 1),
-        PrepareScheme::Current => None,
+        Previous | Both => Some(data_keeper.prev_slider().position() - 1),
+        Current => None,
     };
 
     let current_pos = match scheme {
-        PrepareScheme::Current | PrepareScheme::Both => Some(data_keeper.current_slider().position() - 1),
-        PrepareScheme::Previous => None,
+        Current | Both => Some(data_keeper.current_slider().position() - 1),
+        Previous => None,
     };
 
     let data_positions = DataPositions { prev_pos, current_pos };
 
     let trace_pos = data_keeper.result_states_count();
     data_keeper.new_to_old_pos.insert(trace_pos, data_positions);
-
-    MergerCallResult::CallResult { value, trace_pos }
 }