miden_ace_codegen/layout/
keys.rs1use super::InputLayout;
2use crate::EXT_DEGREE;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub enum InputKey {
7 Public(usize),
9 AuxRandAlpha,
11 AuxRandBeta,
13 MultiAirBetaCore,
16 MultiAirBetaChip,
18 Main {
20 offset: usize,
21 index: usize,
22 },
23 AuxCoord {
25 offset: usize,
26 index: usize,
27 coord: usize,
28 },
29 AuxBusBoundary(usize),
31 VlpiReduction(usize),
33 Gamma,
36 Alpha,
38 ZPowN,
40 ZK,
42 IsFirst,
44 IsLast,
46 IsTransition,
48 IsFirstCore,
52 IsLastCore,
53 IsTransitionCore,
54 IsFirstChip,
56 IsLastChip,
57 IsTransitionChip,
58 Weight0,
60 F,
62 S0,
64 QuotientChunkCoord {
67 offset: usize,
68 chunk: usize,
69 coord: usize,
70 },
71}
72
73#[derive(Debug, Clone, Copy)]
75pub(crate) struct InputKeyMapper<'a> {
76 pub(super) layout: &'a InputLayout,
77}
78
79impl InputKeyMapper<'_> {
80 pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
82 let layout = self.layout;
83 match key {
84 InputKey::Public(i) => layout.regions.public_values.index(i),
85 InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
86 InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
87 InputKey::MultiAirBetaCore => layout.stark.multi_air_beta_core,
88 InputKey::MultiAirBetaChip => layout.stark.multi_air_beta_chip,
89 InputKey::Main { offset, index } => match offset {
90 0 => layout.regions.main_curr.index(index),
91 1 => layout.regions.main_next.index(index),
92 _ => None,
93 },
94 InputKey::AuxCoord { offset, index, coord } => {
95 if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
96 return None;
97 }
98 let local = index * EXT_DEGREE + coord;
99 match offset {
100 0 => layout.regions.aux_curr.index(local),
101 1 => layout.regions.aux_next.index(local),
102 _ => None,
103 }
104 },
105 InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
106 InputKey::VlpiReduction(i) => {
107 let local = i * layout.vlpi_stride;
108 layout.regions.vlpi_reductions.index(local)
109 },
110 InputKey::Alpha => Some(layout.stark.alpha),
112 InputKey::ZPowN => Some(layout.stark.z_pow_n),
113 InputKey::ZK => Some(layout.stark.z_k),
114 InputKey::IsFirst => Some(layout.stark.is_first),
115 InputKey::IsLast => Some(layout.stark.is_last),
116 InputKey::IsTransition => Some(layout.stark.is_transition),
117 InputKey::IsFirstCore => layout.stark.is_first_core,
118 InputKey::IsLastCore => layout.stark.is_last_core,
119 InputKey::IsTransitionCore => layout.stark.is_transition_core,
120 InputKey::IsFirstChip => layout.stark.is_first_chip,
121 InputKey::IsLastChip => layout.stark.is_last_chip,
122 InputKey::IsTransitionChip => layout.stark.is_transition_chip,
123 InputKey::Gamma => Some(layout.stark.gamma),
124 InputKey::Weight0 => Some(layout.stark.weight0),
126 InputKey::F => Some(layout.stark.f),
127 InputKey::S0 => Some(layout.stark.s0),
128 InputKey::QuotientChunkCoord { offset, chunk, coord } => {
129 if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
130 return None;
131 }
132 let idx = chunk * EXT_DEGREE + coord;
133 match offset {
134 0 => layout.regions.quotient_curr.index(idx),
135 1 => layout.regions.quotient_next.index(idx),
136 _ => None,
137 }
138 },
139 }
140 }
141}