miden_ace_codegen/layout/
keys.rs1use super::InputLayout;
2use crate::EXT_DEGREE;
3
4const AIR_SELECTOR_FIRST_OFFSET: usize = 0;
5const AIR_SELECTOR_LAST_OFFSET: usize = 1;
6const AIR_SELECTOR_TRANSITION_OFFSET: usize = 2;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum InputKey {
11 Public(usize),
13 AuxRandAlpha,
15 AuxRandBeta,
17 MultiAirFoldBeta,
19 Preprocessed { offset: usize, index: usize },
21 Main { offset: usize, index: usize },
23 AuxCoord {
25 offset: usize,
26 index: usize,
27 coord: usize,
28 },
29 AuxBusBoundary(usize),
31 Reserved,
33 Alpha,
35 ZPowN,
37 ZK,
40 IsFirst,
42 IsLast,
44 IsTransition,
46 IsFirstAir(usize),
48 IsLastAir(usize),
50 IsTransitionAir(usize),
52 Weight0,
54 F,
56 S0,
58 QuotientChunkCoord {
61 offset: usize,
62 chunk: usize,
63 coord: usize,
64 },
65}
66
67#[derive(Debug, Clone, Copy)]
69pub(crate) struct InputKeyMapper<'a> {
70 pub(super) layout: &'a InputLayout,
71}
72
73impl InputKeyMapper<'_> {
74 pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
76 let layout = self.layout;
77 match key {
78 InputKey::Public(i) => layout.regions.public_values.index(i),
79 InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
80 InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
81 InputKey::MultiAirFoldBeta => layout.stark.multi_air_fold_beta_index(),
82 InputKey::Preprocessed { offset, index } => match offset {
83 0 => layout.regions.preprocessed_curr.index(index),
84 1 => layout.regions.preprocessed_next.index(index),
85 _ => None,
86 },
87 InputKey::Main { offset, index } => match offset {
88 0 => layout.regions.main_curr.index(index),
89 1 => layout.regions.main_next.index(index),
90 _ => None,
91 },
92 InputKey::AuxCoord { offset, index, coord } => {
93 if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
94 return None;
95 }
96 let local = index * EXT_DEGREE + coord;
97 match offset {
98 0 => layout.regions.aux_curr.index(local),
99 1 => layout.regions.aux_next.index(local),
100 _ => None,
101 }
102 },
103 InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
104 InputKey::Reserved => Some(layout.stark.reserved),
105 InputKey::Alpha => Some(layout.stark.alpha),
106 InputKey::ZPowN => Some(layout.stark.z_pow_n),
107 InputKey::ZK => Some(layout.stark.z_k),
108 InputKey::IsFirst => Some(layout.stark.is_first),
109 InputKey::IsLast => Some(layout.stark.is_last),
110 InputKey::IsTransition => Some(layout.stark.is_transition),
111 InputKey::IsFirstAir(i) => {
112 layout.stark.air_selector_index(i, AIR_SELECTOR_FIRST_OFFSET)
113 },
114 InputKey::IsLastAir(i) => layout.stark.air_selector_index(i, AIR_SELECTOR_LAST_OFFSET),
115 InputKey::IsTransitionAir(i) => {
116 layout.stark.air_selector_index(i, AIR_SELECTOR_TRANSITION_OFFSET)
117 },
118 InputKey::Weight0 => Some(layout.stark.weight0),
119 InputKey::F => Some(layout.stark.f),
120 InputKey::S0 => Some(layout.stark.s0),
121 InputKey::QuotientChunkCoord { offset, chunk, coord } => {
122 if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
123 return None;
124 }
125 let idx = chunk * EXT_DEGREE + coord;
126 match offset {
127 0 => layout.regions.quotient_curr.index(idx),
128 1 => layout.regions.quotient_next.index(idx),
129 _ => None,
130 }
131 },
132 }
133 }
134}