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 MultiAirFoldCoeff(usize),
26 Preprocessed { offset: usize, index: usize },
28 Main { offset: usize, index: usize },
30 AuxCoord {
32 offset: usize,
33 index: usize,
34 coord: usize,
35 },
36 AuxBusBoundary(usize),
38 Reserved,
40 Alpha,
42 ZPowN,
44 ZK,
47 IsFirst,
49 IsLast,
51 IsTransition,
53 IsFirstAir(usize),
55 IsLastAir(usize),
57 IsTransitionAir(usize),
59 Weight0,
61 F,
64 S0,
66 QuotientChunkCoord {
69 offset: usize,
70 chunk: usize,
71 coord: usize,
72 },
73}
74
75#[derive(Debug, Clone, Copy)]
77pub(crate) struct InputKeyMapper<'a> {
78 pub(super) layout: &'a InputLayout,
79}
80
81impl InputKeyMapper<'_> {
82 pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
84 let layout = self.layout;
85 match key {
86 InputKey::Public(i) => layout.regions.public_values.index(i),
87 InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
88 InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
89 InputKey::MultiAirFoldBeta => layout.stark.multi_air_fold_beta_index(),
90 InputKey::MultiAirFoldCoeff(_) => None,
92 InputKey::Preprocessed { offset, index } => match offset {
93 0 => layout.regions.preprocessed_curr.index(index),
94 1 => layout.regions.preprocessed_next.index(index),
95 _ => None,
96 },
97 InputKey::Main { offset, index } => match offset {
98 0 => layout.regions.main_curr.index(index),
99 1 => layout.regions.main_next.index(index),
100 _ => None,
101 },
102 InputKey::AuxCoord { offset, index, coord } => {
103 if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
104 return None;
105 }
106 let local = index * EXT_DEGREE + coord;
107 match offset {
108 0 => layout.regions.aux_curr.index(local),
109 1 => layout.regions.aux_next.index(local),
110 _ => None,
111 }
112 },
113 InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
114 InputKey::Reserved => Some(layout.stark.reserved),
115 InputKey::Alpha => Some(layout.stark.alpha),
116 InputKey::ZPowN => Some(layout.stark.z_pow_n),
117 InputKey::ZK => Some(layout.stark.z_k),
118 InputKey::IsFirst => Some(layout.stark.is_first),
119 InputKey::IsLast => Some(layout.stark.is_last),
120 InputKey::IsTransition => Some(layout.stark.is_transition),
121 InputKey::IsFirstAir(i) => {
122 layout.stark.air_selector_index(i, AIR_SELECTOR_FIRST_OFFSET)
123 },
124 InputKey::IsLastAir(i) => layout.stark.air_selector_index(i, AIR_SELECTOR_LAST_OFFSET),
125 InputKey::IsTransitionAir(i) => {
126 layout.stark.air_selector_index(i, AIR_SELECTOR_TRANSITION_OFFSET)
127 },
128 InputKey::Weight0 => Some(layout.stark.weight0),
129 InputKey::F => Some(layout.stark.f),
130 InputKey::S0 => Some(layout.stark.s0),
131 InputKey::QuotientChunkCoord { offset, chunk, coord } => {
132 if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
133 return None;
134 }
135 let idx = chunk * EXT_DEGREE + coord;
136 match offset {
137 0 => layout.regions.quotient_curr.index(idx),
138 1 => layout.regions.quotient_next.index(idx),
139 _ => None,
140 }
141 },
142 }
143 }
144}