1  // Copyright (c) 2019-2025 Alpha-Delta Network Inc.
  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  mod equal;
 17  mod from;
 18  mod from_private_key;
 19  mod helpers;
 20  mod ternary;
 21  mod to_address;
 22  
 23  #[cfg(test)]
 24  use alphavm_circuit_types::environment::{assert_count, assert_output_mode, assert_scope};
 25  
 26  use crate::PrivateKey;
 27  use alphavm_circuit_network::Alpha;
 28  use alphavm_circuit_types::{Address, Boolean, Field, Group, Scalar, environment::prelude::*};
 29  
 30  #[derive(Clone)]
 31  pub struct ComputeKey<A: Alpha> {
 32      /// The signature public key `pk_sig` := G^sk_sig.
 33      pk_sig: Group<A>,
 34      /// The signature public randomizer `pr_sig` := G^r_sig.
 35      pr_sig: Group<A>,
 36      /// The PRF secret key `sk_prf` := RO(G^sk_sig || G^r_sig).
 37      sk_prf: Scalar<A>,
 38  }
 39  
 40  impl<A: Alpha> Inject for ComputeKey<A> {
 41      type Primitive = console::ComputeKey<A::Network>;
 42  
 43      /// Initializes an account compute key from the given mode and native compute key.
 44      fn new(mode: Mode, compute_key: Self::Primitive) -> Self {
 45          // Inject `pk_sig`.
 46          let pk_sig = Group::new(mode, compute_key.pk_sig());
 47          // Inject `pr_sig`.
 48          let pr_sig = Group::new(mode, compute_key.pr_sig());
 49          // Output the compute key.
 50          Self::from((pk_sig, pr_sig))
 51      }
 52  }
 53  
 54  impl<A: Alpha> ComputeKey<A> {
 55      /// Returns the signature public key.
 56      pub const fn pk_sig(&self) -> &Group<A> {
 57          &self.pk_sig
 58      }
 59  
 60      /// Returns the signature public randomizer.
 61      pub const fn pr_sig(&self) -> &Group<A> {
 62          &self.pr_sig
 63      }
 64  
 65      /// Returns the PRF secret key.
 66      pub const fn sk_prf(&self) -> &Scalar<A> {
 67          &self.sk_prf
 68      }
 69  }
 70  
 71  impl<A: Alpha> Eject for ComputeKey<A> {
 72      type Primitive = console::ComputeKey<A::Network>;
 73  
 74      /// Ejects the mode of the compute key.
 75      fn eject_mode(&self) -> Mode {
 76          (&self.pk_sig, &self.pr_sig, &self.sk_prf).eject_mode()
 77      }
 78  
 79      /// Ejects the compute key.
 80      fn eject_value(&self) -> Self::Primitive {
 81          match Self::Primitive::try_from((&self.pk_sig, &self.pr_sig).eject_value()) {
 82              Ok(compute_key) => compute_key,
 83              Err(error) => A::halt(format!("Failed to eject the compute key: {error}")),
 84          }
 85      }
 86  }
 87  
 88  #[cfg(test)]
 89  pub(crate) mod tests {
 90      use super::*;
 91      use crate::{Circuit, helpers::generate_account};
 92  
 93      use anyhow::Result;
 94  
 95      const ITERATIONS: u64 = 250;
 96  
 97      fn check_new(
 98          mode: Mode,
 99          num_constants: u64,
100          num_public: u64,
101          num_private: u64,
102          num_constraints: u64,
103      ) -> Result<()> {
104          for i in 0..ITERATIONS {
105              // Generate a private key, compute key, view key, and address.
106              let (_private_key, compute_key, _view_key, _address) = generate_account()?;
107  
108              Circuit::scope(format!("New {mode}"), || {
109                  let candidate = ComputeKey::<Circuit>::new(mode, compute_key);
110                  match mode.is_constant() {
111                      true => assert_eq!(Mode::Constant, candidate.eject_mode()),
112                      false => assert_eq!(Mode::Private, candidate.eject_mode()),
113                  };
114                  assert_eq!(compute_key, candidate.eject_value());
115                  // TODO (howardwu): Resolve skipping the cost count checks for the burn-in round.
116                  if i > 0 {
117                      assert_scope!(num_constants, num_public, num_private, num_constraints);
118                  }
119              });
120              Circuit::reset();
121          }
122          Ok(())
123      }
124  
125      #[test]
126      fn test_compute_key_new_constant() -> Result<()> {
127          check_new(Mode::Constant, 274, 0, 0, 0)
128      }
129  
130      #[test]
131      fn test_compute_key_new_public() -> Result<()> {
132          check_new(Mode::Public, 9, 4, 869, 873)
133      }
134  
135      #[test]
136      fn test_compute_key_new_private() -> Result<()> {
137          check_new(Mode::Private, 9, 0, 873, 873)
138      }
139  }