refactor(aquavm): improve readability of ap merger and handler (#354)

At the moment `Ap` state contains vector of generations, although the ap instruction itself allow to specify only one stream. It was done to support several streams in a result position in a future. But the current realization allows to handle empty vector that was implemented to support states for scalars and it turned out (for more info see #326) that states for scalars aren't actually needed.

Closes #355.

crates/air-lib/interpreter-data/src/executed_state/impls.rs

@@ -79,9 +79,9 @@ impl ExecutedState {
 }
 
 impl ApResult {
-    pub fn new(res_gens: Vec<u32>) -> Self {
+    pub fn new(res_generation: u32) -> Self {
         Self {
-            res_generations: res_gens,
+            res_generations: vec![res_generation],
         }
     }
 }

crates/air-lib/test-utils/src/executed_state.rs

@@ -132,10 +132,8 @@ pub fn subtrace_desc(begin_pos: impl Into<TracePos>, subtrace_len: u32) -> SubTr
     }
 }
 
-pub fn ap(dst: Option<u32>) -> ExecutedState {
-    let res_generations = option_to_vec(dst);
-    let ap_result = ApResult::new(res_generations);
-
+pub fn ap(generation: u32) -> ExecutedState {
+    let ap_result = ApResult::new(generation);
     ExecutedState::Ap(ap_result)
 }
 
@@ -143,10 +141,3 @@ pub fn canon(stream_elements_pos: Vec<TracePos>) -> ExecutedState {
     let canon_result = CanonResult::new(stream_elements_pos);
     ExecutedState::Canon(canon_result)
 }
-
-fn option_to_vec(maybe_value: Option<u32>) -> Vec<u32> {
-    match maybe_value {
-        Some(value) => vec![value],
-        None => vec![],
-    }
-}

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

@@ -21,11 +21,17 @@ const EXPECTED_STATE_NAME: &str = "ap";
 #[derive(Debug, Clone)]
 pub enum MergerApResult {
     /// There is no corresponding state in a trace for this call.
-    Empty,
+    NotMet,
 
     /// There was a state in at least one of the contexts. If there were two states in
     /// both contexts, they were successfully merged.
-    ApResult { res_generation: Option<u32> },
+    Met(MetApResult),
+}
+
+#[derive(Debug, Clone)]
+pub struct MetApResult {
+    pub generation: u32,
+    pub value_source: ValueSource,
 }
 
 pub(crate) fn try_merge_next_state_as_ap(data_keeper: &mut DataKeeper) -> MergeResult<MergerApResult> {
@@ -36,12 +42,12 @@ pub(crate) fn try_merge_next_state_as_ap(data_keeper: &mut DataKeeper) -> MergeR
     let current_state = data_keeper.current_slider_mut().next_state();
 
     match (prev_state, current_state) {
-        (Some(Ap(prev_ap)), Some(Ap(_))) => prepare_merge_result(Some(prev_ap), Both, data_keeper),
-        (Some(Ap(prev_ap)), None) => prepare_merge_result(Some(prev_ap), Previous, data_keeper),
+        (Some(Ap(prev_ap)), Some(Ap(_))) => prepare_merge_result(prev_ap, Both, data_keeper),
+        (Some(Ap(prev_ap)), None) => prepare_merge_result(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(_))) => prepare_merge_result(None, Current, data_keeper),
-        (None, None) => Ok(MergerApResult::Empty),
+        (None, Some(Ap(current_ap))) => prepare_merge_result(current_ap, Current, data_keeper),
+        (None, None) => Ok(MergerApResult::NotMet),
         (prev_state, current_state) => Err(MergeError::incompatible_states(
             prev_state,
             current_state,
@@ -50,24 +56,10 @@ pub(crate) fn try_merge_next_state_as_ap(data_keeper: &mut DataKeeper) -> MergeR
     }
 }
 
-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 {
     ($ap_result:ident, $generations:expr, $error_ty:ident) => {
         match $generations.len() {
-            0 => None,
-            1 => Some($generations[0]),
+            1 => $generations[0],
             _ => {
                 let ap_error = super::ApResultError::$error_ty($ap_result);
                 return Err(super::MergeError::IncorrectApResult(ap_error));
@@ -76,10 +68,25 @@ macro_rules! to_maybe_generation {
     };
 }
 
-fn to_merger_result(ap_result: ApResult) -> MergeResult<MergerApResult> {
-    let res_generation = to_maybe_generation!(ap_result, &ap_result.res_generations, TooManyDstGenerations);
+fn prepare_merge_result(
+    ap_result: ApResult,
+    scheme: PreparationScheme,
+    data_keeper: &mut DataKeeper,
+) -> MergeResult<MergerApResult> {
+    prepare_positions_mapping(scheme, data_keeper);
 
-    let ap_result = MergerApResult::ApResult { res_generation };
+    let generation = to_maybe_generation!(ap_result, &ap_result.res_generations, InvalidDstGenerations);
+    let met_result = MetApResult::new(generation, scheme.into());
+    let ap_result = MergerApResult::Met(met_result);
 
     Ok(ap_result)
 }
+
+impl MetApResult {
+    pub(crate) fn new(generation: u32, value_source: ValueSource) -> Self {
+        Self {
+            generation,
+            value_source,
+        }
+    }
+}

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

@@ -29,22 +29,16 @@ pub enum MergerCallResult {
 
     /// There was a state in at least one of the contexts. If there were two states in
     /// both contexts, they were successfully merged.
-    Met(MetResult),
+    Met(MetCallResult),
 }
 
 #[derive(Debug, Clone)]
-pub struct MetResult {
+pub struct MetCallResult {
     pub result: CallResult,
     pub trace_pos: TracePos,
     pub source: ValueSource,
 }
 
-#[derive(Debug, Clone, Copy)]
-pub enum ValueSource {
-    PreviousData,
-    CurrentData,
-}
-
 pub(crate) fn try_merge_next_state_as_call(data_keeper: &mut DataKeeper) -> MergeResult<MergerCallResult> {
     use ExecutedState::Call;
     use PreparationScheme::*;
@@ -99,7 +93,7 @@ pub(super) fn prepare_call_result(
     let trace_pos = data_keeper.result_trace_next_pos();
     prepare_positions_mapping(scheme, data_keeper);
 
-    let met_result = MetResult::new(call_result, trace_pos, scheme.into());
+    let met_result = MetCallResult::new(call_result, trace_pos, scheme.into());
     MergerCallResult::Met(met_result)
 }
 
@@ -112,7 +106,7 @@ impl From<PreparationScheme> for ValueSource {
     }
 }
 
-impl MetResult {
+impl MetCallResult {
     pub fn new(result: CallResult, trace_pos: TracePos, source: ValueSource) -> Self {
         Self {
             result,

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

@@ -55,9 +55,9 @@ pub enum MergeError {
 
 #[derive(ThisError, Debug)]
 pub enum ApResultError {
-    /// Error occurred when Ap results contains more then 1 generation in destination.
-    #[error("{0:?} ap result contains too many generations in destination")]
-    TooManyDstGenerations(ApResult),
+    /// Error occurred when Ap results contains not 1 generation in destination.
+    #[error("{0:?} ap result contains inappropriate generation count in destination")]
+    InvalidDstGenerations(ApResult),
 }
 
 #[derive(ThisError, Debug)]

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

@@ -240,7 +240,7 @@ mod tests {
 
         let fold_result = FoldResult { lore };
 
-        let slider = TraceSlider::new(vec![ExecutedState::Ap(ApResult::new(vec![0]))]);
+        let slider = TraceSlider::new(vec![ExecutedState::Ap(ApResult::new(0))]);
         let ctx = MergeCtx { slider };
 
         let (all_states, convoluted_lens) =
@@ -285,9 +285,9 @@ mod tests {
         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])),
+            ExecutedState::Ap(ApResult::new(0)),
+            ExecutedState::Ap(ApResult::new(1)),
+            ExecutedState::Ap(ApResult::new(2)),
         ]);
         let ctx = MergeCtx { slider };
 

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

@@ -23,9 +23,9 @@ mod par_merger;
 mod position_mapping;
 
 pub use ap_merger::MergerApResult;
+pub use ap_merger::MetApResult;
 pub use call_merger::MergerCallResult;
-pub use call_merger::MetResult;
-pub use call_merger::ValueSource;
+pub use call_merger::MetCallResult;
 pub use canon_merger::MergerCanonResult;
 pub use fold_merger::MergerFoldResult;
 pub use par_merger::MergerParResult;
@@ -60,6 +60,12 @@ pub enum MergeCtxType {
     Previous,
 }
 
+#[derive(Debug, Clone, Copy)]
+pub enum ValueSource {
+    PreviousData,
+    CurrentData,
+}
+
 use std::fmt;
 
 impl fmt::Display for MergeCtxType {