feat: add methods to get air widths/interactions from vkey (openvm-org#79)

* feat: add fns for getting widths/interactions from vkey

* feat(test): add test for vkey widths/interactions method

* chore: fix clippy

Author: shuklaayush

crates/stark-backend/src/keygen/types.rs

@@ -183,6 +183,14 @@ impl<SC: StarkGenericConfig> MultiStarkVerifyingKey<SC> {
     pub fn num_challenges_per_phase(&self) -> Vec<usize> {
         self.full_view().num_challenges_per_phase()
     }
+
+    pub fn total_widths(&self) -> Vec<usize> {
+        self.full_view().total_widths::<SC::Challenge>()
+    }
+
+    pub fn num_interactions(&self) -> Vec<usize> {
+        self.full_view().num_interactions()
+    }
 }
 
 /// Prover only data for preprocessed trace for a single AIR.

crates/stark-backend/src/keygen/view.rs

@@ -1,4 +1,5 @@
 use itertools::Itertools;
+use p3_field::{ExtensionField, Field};
 
 use crate::{
     config::{Com, StarkGenericConfig, Val},
@@ -75,4 +76,29 @@ impl<Val, Com: Clone> MultiStarkVerifyingKeyView<'_, Val, Com> {
             .max()
             .unwrap_or_else(|| panic!("No challenges used in challenge phase {phase_idx}"))
     }
+
+    /// Returns the total width for each AIR.
+    pub fn total_widths<E>(&self) -> Vec<usize>
+    where
+        Val: Field,
+        E: ExtensionField<Val>,
+    {
+        self.per_air
+            .iter()
+            .map(|vk| {
+                vk.params.width.preprocessed.unwrap_or(0)
+                    + vk.params.width.cached_mains.iter().sum::<usize>()
+                    + vk.params.width.common_main
+                    + vk.params.width.after_challenge.iter().sum::<usize>() * E::D
+            })
+            .collect()
+    }
+
+    /// Returns the number of interactions for each AIR.
+    pub fn num_interactions(&self) -> Vec<usize> {
+        self.per_air
+            .iter()
+            .map(|vk| vk.symbolic_constraints.interactions.len())
+            .collect()
+    }
 }

crates/stark-backend/tests/integration_test.rs

@@ -226,6 +226,45 @@ fn test_optional_air() {
     }
 }
 
+#[test]
+fn test_vkey_methods() {
+    use openvm_stark_backend::engine::StarkEngine;
+    use openvm_stark_sdk::{
+        config::{setup_tracing, FriParameters},
+        dummy_airs::fib_air::air::FibonacciAir,
+    };
+    use p3_air::BaseAir;
+
+    setup_tracing();
+
+    let engine = BabyBearPoseidon2Engine::new(FriParameters::standard_fast());
+
+    let fib_chip = FibonacciChip::new(0, 1, 8);
+    let send_chip = DummyInteractionChip::new_without_partition(1, true, 0);
+    let recv_chip = DummyInteractionChip::new_without_partition(1, false, 0);
+
+    let mut keygen_builder = engine.keygen_builder();
+    let _ = keygen_builder.add_air(fib_chip.air());
+    let _ = keygen_builder.add_air(send_chip.air());
+    let _ = keygen_builder.add_air(recv_chip.air());
+    let pk = keygen_builder.generate_pk();
+
+    let vk = pk.get_vk();
+
+    let widths = vk.total_widths();
+    assert_eq!(widths.len(), 3);
+    assert_eq!(widths[0], BaseAir::<BabyBear>::width(&FibonacciAir));
+    // 1 interaction -> 1 perm col + 1 cumsum col -> 8 cols in base field
+    assert_eq!(widths[1], BaseAir::<BabyBear>::width(&send_chip.air) + 8);
+    assert_eq!(widths[2], BaseAir::<BabyBear>::width(&recv_chip.air) + 8);
+
+    let interactions = vk.num_interactions();
+    assert_eq!(interactions.len(), 3);
+    assert_eq!(interactions[0], 0);
+    assert_eq!(interactions[1], 1);
+    assert_eq!(interactions[2], 1);
+}
+
 fn get_fib_number(n: usize) -> u32 {
     let mut a = 0;
     let mut b = 1;