1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaVM library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use super::*;
 17  
 18  impl<E: Environment> Ternary for Boolean<E> {
 19      type Boolean = Boolean<E>;
 20      type Output = Self;
 21  
 22      /// Returns `first` if `condition` is `true`, otherwise returns `second`.
 23      fn ternary(condition: &Self::Boolean, first: &Self, second: &Self) -> Self::Output {
 24          // Constant `condition`
 25          if condition.is_constant() {
 26              match condition.eject_value() {
 27                  true => first.clone(),
 28                  false => second.clone(),
 29              }
 30          }
 31          // Constant `first`
 32          else if first.is_constant() {
 33              match first.eject_value() {
 34                  true => condition | second,
 35                  false => !condition & second,
 36              }
 37          }
 38          // Constant `second`
 39          else if second.is_constant() {
 40              match second.eject_value() {
 41                  true => !condition | first,
 42                  false => condition & first,
 43              }
 44          }
 45          // Variables
 46          else {
 47              // Compute the witness value, based on the condition.
 48              let witness = match condition.eject_value() {
 49                  true => first.eject_value(),
 50                  false => second.eject_value(),
 51              };
 52  
 53              // Declare a new variable with the expected output as witness.
 54              // Note: The constraint below will ensure `output` is either 0 or 1,
 55              // assuming `self` and `other` are well-formed (they are either 0 or 1).
 56              let output = Boolean(
 57                  E::new_variable(
 58                      Mode::Private,
 59                      match witness {
 60                          true => E::BaseField::one(),
 61                          false => E::BaseField::zero(),
 62                      },
 63                  )
 64                  .into(),
 65              );
 66  
 67              //
 68              // Ternary Enforcement
 69              // -------------------------------------------------------
 70              //    output = condition * a + (1 - condition) * b
 71              // => output = b + condition * (a - b)
 72              // => condition * (a - b) = output - b
 73              //
 74              // See `Field::ternary()` for the proof of correctness.
 75              //
 76              E::enforce(|| (condition, (&first.0 - &second.0), (&output.0 - &second.0)));
 77  
 78              output
 79          }
 80      }
 81  }
 82  
 83  #[cfg(test)]
 84  mod tests {
 85      use super::*;
 86      use alphavm_circuit_environment::Circuit;
 87  
 88      fn check_ternary(
 89          name: &str,
 90          expected: bool,
 91          condition: Boolean<Circuit>,
 92          a: Boolean<Circuit>,
 93          b: Boolean<Circuit>,
 94          num_constants: u64,
 95          num_public: u64,
 96          num_private: u64,
 97          num_constraints: u64,
 98      ) {
 99          Circuit::scope(name, || {
100              let case = format!("({} ? {} : {})", condition.eject_value(), a.eject_value(), b.eject_value());
101              let candidate = Boolean::ternary(&condition, &a, &b);
102              assert_eq!(expected, candidate.eject_value(), "{case}");
103              assert_scope!(num_constants, num_public, num_private, num_constraints);
104          });
105      }
106  
107      fn run_test(
108          mode_condition: Mode,
109          mode_a: Mode,
110          mode_b: Mode,
111          num_constants: u64,
112          num_public: u64,
113          num_private: u64,
114          num_constraints: u64,
115      ) {
116          for flag in [true, false] {
117              for first in [true, false] {
118                  for second in [true, false] {
119                      let condition = Boolean::<Circuit>::new(mode_condition, flag);
120                      let a = Boolean::<Circuit>::new(mode_a, first);
121                      let b = Boolean::<Circuit>::new(mode_b, second);
122  
123                      let name = format!("{mode_condition} ? {mode_a} : {mode_b}");
124                      check_ternary(
125                          &name,
126                          if flag { first } else { second },
127                          condition,
128                          a,
129                          b,
130                          num_constants,
131                          num_public,
132                          num_private,
133                          num_constraints,
134                      );
135                  }
136              }
137          }
138      }
139  
140      #[test]
141      fn test_if_constant_then_constant_else_constant() {
142          run_test(Mode::Constant, Mode::Constant, Mode::Constant, 0, 0, 0, 0);
143      }
144  
145      #[test]
146      fn test_if_constant_then_constant_else_public() {
147          run_test(Mode::Constant, Mode::Constant, Mode::Public, 0, 0, 0, 0);
148      }
149  
150      #[test]
151      fn test_if_constant_then_constant_else_private() {
152          run_test(Mode::Constant, Mode::Constant, Mode::Private, 0, 0, 0, 0);
153      }
154  
155      #[test]
156      fn test_if_constant_then_public_else_constant() {
157          run_test(Mode::Constant, Mode::Public, Mode::Constant, 0, 0, 0, 0);
158      }
159  
160      #[test]
161      fn test_if_constant_then_public_else_public() {
162          run_test(Mode::Constant, Mode::Public, Mode::Public, 0, 0, 0, 0);
163      }
164  
165      #[test]
166      fn test_if_constant_then_public_else_private() {
167          run_test(Mode::Constant, Mode::Public, Mode::Private, 0, 0, 0, 0);
168      }
169  
170      #[test]
171      fn test_if_constant_then_private_else_constant() {
172          run_test(Mode::Constant, Mode::Private, Mode::Constant, 0, 0, 0, 0);
173      }
174  
175      #[test]
176      fn test_if_constant_then_private_else_public() {
177          run_test(Mode::Constant, Mode::Private, Mode::Public, 0, 0, 0, 0);
178      }
179  
180      #[test]
181      fn test_if_constant_then_private_else_private() {
182          run_test(Mode::Constant, Mode::Private, Mode::Private, 0, 0, 0, 0);
183      }
184  
185      #[test]
186      fn test_if_public_then_constant_else_constant() {
187          run_test(Mode::Public, Mode::Constant, Mode::Constant, 0, 0, 0, 0);
188      }
189  
190      #[test]
191      fn test_if_public_then_constant_else_public() {
192          run_test(Mode::Public, Mode::Constant, Mode::Public, 0, 0, 1, 1);
193      }
194  
195      #[test]
196      fn test_if_public_then_constant_else_private() {
197          run_test(Mode::Public, Mode::Constant, Mode::Private, 0, 0, 1, 1);
198      }
199  
200      #[test]
201      fn test_if_public_then_public_else_constant() {
202          run_test(Mode::Public, Mode::Public, Mode::Constant, 0, 0, 1, 1);
203      }
204  
205      #[test]
206      fn test_if_public_then_public_else_public() {
207          run_test(Mode::Public, Mode::Public, Mode::Public, 0, 0, 1, 1);
208      }
209  
210      #[test]
211      fn test_if_public_then_public_else_private() {
212          run_test(Mode::Public, Mode::Public, Mode::Private, 0, 0, 1, 1);
213      }
214  
215      #[test]
216      fn test_if_public_then_private_else_constant() {
217          run_test(Mode::Public, Mode::Private, Mode::Constant, 0, 0, 1, 1);
218      }
219  
220      #[test]
221      fn test_if_public_then_private_else_public() {
222          run_test(Mode::Public, Mode::Private, Mode::Public, 0, 0, 1, 1);
223      }
224  
225      #[test]
226      fn test_if_public_then_private_else_private() {
227          run_test(Mode::Public, Mode::Private, Mode::Private, 0, 0, 1, 1);
228      }
229  
230      #[test]
231      fn test_if_private_then_constant_else_constant() {
232          run_test(Mode::Private, Mode::Constant, Mode::Constant, 0, 0, 0, 0);
233      }
234  
235      #[test]
236      fn test_if_private_then_constant_else_public() {
237          run_test(Mode::Private, Mode::Constant, Mode::Public, 0, 0, 1, 1);
238      }
239  
240      #[test]
241      fn test_if_private_then_constant_else_private() {
242          run_test(Mode::Private, Mode::Constant, Mode::Private, 0, 0, 1, 1);
243      }
244  
245      #[test]
246      fn test_if_private_then_public_else_constant() {
247          run_test(Mode::Private, Mode::Public, Mode::Constant, 0, 0, 1, 1);
248      }
249  
250      #[test]
251      fn test_if_private_then_public_else_public() {
252          run_test(Mode::Private, Mode::Public, Mode::Public, 0, 0, 1, 1);
253      }
254  
255      #[test]
256      fn test_if_private_then_public_else_private() {
257          run_test(Mode::Private, Mode::Public, Mode::Private, 0, 0, 1, 1);
258      }
259  
260      #[test]
261      fn test_if_private_then_private_else_constant() {
262          run_test(Mode::Private, Mode::Private, Mode::Constant, 0, 0, 1, 1);
263      }
264  
265      #[test]
266      fn test_if_private_then_private_else_public() {
267          run_test(Mode::Private, Mode::Private, Mode::Public, 0, 0, 1, 1);
268      }
269  
270      #[test]
271      fn test_if_private_then_private_else_private() {
272          run_test(Mode::Private, Mode::Private, Mode::Private, 0, 0, 1, 1);
273      }
274  }