Fix ap trace handler behaviour (#224)

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 }
 }