chore(sdk): docs, cleanup (#14)

Makefile

@@ -1,8 +1,9 @@
 .PHONY: deps clean example
-
 deps:
 	@cargo install cross --git https://github.com/cross-rs/cross.git --rev 1511a28
+
 clean:
 	@cargo clean
+
 example:
 	@cargo run --release -p stealth_address_kit_example

sdk/src/baby_jub_jub_impl.rs

@@ -1,6 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
-
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the BabyJubJub curve.
 use ark_ed_on_bn254::{EdwardsProjective, Fr};
 
 pub struct BabyJubJub;

sdk/src/bls12_377_impl.rs

@@ -1,6 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
-
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Bls12_377 curve.
 use ark_bls12_377::{Bls12_377, Fr, G1Projective};
 
 impl StealthAddressOnCurve for Bls12_377 {
@@ -12,4 +11,4 @@ impl StealthAddressOnCurve for Bls12_377 {
 use crate::define_curve_ffi;
 #[cfg(feature = "ffi")]
 define_curve_ffi!(bls12_377, Bls12_377, Fr, G1Projective, 32, 48);
-define_curve_tests!(ark_bls12_377::Bls12_377);
+define_curve_tests!(Bls12_377);

sdk/src/bls12_381_impl.rs

@@ -1,6 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
-
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Bls12_381 curve.
 use ark_bls12_381::{Bls12_381, Fr, G1Projective};
 
 impl StealthAddressOnCurve for Bls12_381 {

sdk/src/bn254_impl.rs

@@ -1,5 +1,4 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
 use ark_bn254::{Bn254, Fr, G1Projective};
 
 impl StealthAddressOnCurve for Bn254 {

sdk/src/bw6_761_impl.rs

@@ -1,6 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
-
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the BW6_761 curve.
 use ark_bw6_761::{Fr, G1Projective, BW6_761};
 
 impl StealthAddressOnCurve for BW6_761 {

sdk/src/pallas_impl.rs

@@ -1,5 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Pallas curve.
 use ark_pallas::{Fr, Projective};
 
 pub struct Pallas;

sdk/src/secp256k1_impl.rs

@@ -1,5 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Secp256k1 curve.
 use ark_secp256k1::{Fr, Projective};
 
 pub struct Secp256k1;

sdk/src/secp256r1_impl.rs

@@ -1,5 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Secp256r1 curve.
 use ark_secp256r1::{Fr, Projective};
 
 pub struct Secp256r1;

sdk/src/stealth_addresses.rs

@@ -5,58 +5,92 @@ use ark_std::rand::rngs::OsRng;
 use ark_std::UniformRand;
 use std::fmt::Display;
 use std::ops::{Add, Mul};
-
 use tiny_keccak::{Hasher, Keccak};
 
+/// A trait for types that have a view tag.
 pub trait HasViewTag {
+    /// Returns the view tag.
     fn get_view_tag(&self) -> u64;
 }
 
-// Implement HasViewTag for any Fp type
 impl<P: FpConfig<N>, const N: usize> HasViewTag for Fp<P, N> {
     fn get_view_tag(&self) -> u64 {
         self.0 .0[0]
     }
 }
 
+/// A trait for converting projective points to bytes.
 pub trait ToBytesFromProjective {
+    /// Converts the projective point to a byte vector.
     fn to_bytes(&self) -> Vec<u8>;
 }
 
-// Implement ToBytesFromProjective for any ProjectiveCurve
 impl<G: CurveGroup> ToBytesFromProjective for G
 where
     G::Affine: CanonicalSerialize,
 {
     fn to_bytes(&self) -> Vec<u8> {
         let affine = self.into_affine();
-        let mut bytes = Vec::new();
+        let mut bytes = Vec::with_capacity(affine.compressed_size());
         affine.serialize_compressed(&mut bytes).unwrap();
         bytes
     }
 }
 
+/// A trait for implementing stealth addresses on elliptic curves.
 pub trait StealthAddressOnCurve {
+    /// The projective representation of the elliptic curve point.
     type Projective: Display
         + Add<Output = Self::Projective>
         + Mul<Self::Fr, Output = Self::Projective>
         + From<<Self::Projective as CurveGroup>::Affine>
-        + ark_ec::CurveGroup;
-    type Fr: Add<Self::Fr, Output = Self::Fr> + ark_ff::PrimeField + HasViewTag;
+        + CurveGroup;
 
+    /// The scalar field of the elliptic curve.
+    type Fr: Add<Self::Fr, Output = Self::Fr> + PrimeField + HasViewTag;
+
+    /// Derives a public key from a given private key.
+    ///
+    /// # Arguments
+    ///
+    /// * `private_key` - A reference to the private key.
+    ///
+    /// # Returns
+    ///
+    /// The derived public key.
     fn derive_public_key(private_key: &Self::Fr) -> Self::Projective {
         Self::Projective::generator() * *private_key
     }
 
+    /// Generates a random keypair.
+    ///
+    /// # Returns
+    ///
+    /// A tuple containing the private key and the derived public key.
     fn random_keypair() -> (Self::Fr, Self::Projective) {
         let private_key = Self::generate_random_fr();
         let public_key = Self::derive_public_key(&private_key);
         (private_key, public_key)
     }
+
+    /// Generates a random scalar field element.
+    ///
+    /// # Returns
+    ///
+    /// A random scalar field element.
     fn generate_random_fr() -> Self::Fr {
-        let mut rng = OsRng;
-        Self::Fr::rand(&mut rng)
+        Self::Fr::rand(&mut OsRng)
     }
+
+    /// Hashes an input byte slice to a scalar field element.
+    ///
+    /// # Arguments
+    ///
+    /// * `input` - A byte slice to be hashed.
+    ///
+    /// # Returns
+    ///
+    /// A scalar field element derived from the hash of the input.
     fn hash_to_fr(input: &[u8]) -> Self::Fr {
         let mut hash = [0; 32];
         let mut hasher = Keccak::v256();
@@ -66,6 +100,17 @@ pub trait StealthAddressOnCurve {
         // We export the hash as a field element
         Self::Fr::from_le_bytes_mod_order(hash.as_slice())
     }
+
+    /// Computes a shared elliptic curve point given a private key and a public key.
+    ///
+    /// # Arguments
+    ///
+    /// * `private_key` - The private key.
+    /// * `public_key` - The public key.
+    ///
+    /// # Returns
+    ///
+    /// The computed shared elliptic curve point.
     fn compute_shared_point(
         private_key: Self::Fr,
         public_key: Self::Projective,
@@ -73,6 +118,17 @@ pub trait StealthAddressOnCurve {
         public_key * private_key
     }
 
+    /// Generates a stealth address.
+    ///
+    /// # Arguments
+    ///
+    /// * `viewing_public_key` - The viewing public key.
+    /// * `spending_public_key` - The spending public key.
+    /// * `ephemeral_private_key` - The ephemeral private key.
+    ///
+    /// # Returns
+    ///
+    /// A tuple containing the stealth address and the view tag.
     fn generate_stealth_address(
         viewing_public_key: Self::Projective,
         spending_public_key: Self::Projective,
@@ -85,6 +141,18 @@ pub trait StealthAddressOnCurve {
         (q_hashed_in_g1 + spending_public_key, view_tag)
     }
 
+    /// Generates a stealth private key.
+    ///
+    /// # Arguments
+    ///
+    /// * `ephemeral_public_key` - The ephemeral public key.
+    /// * `viewing_key` - The viewing key.
+    /// * `spending_key` - The spending key.
+    /// * `expected_view_tag` - The expected view tag.
+    ///
+    /// # Returns
+    ///
+    /// An optional stealth private key.
     fn generate_stealth_private_key(
         ephemeral_public_key: Self::Projective,
         viewing_key: Self::Fr,
@@ -92,14 +160,9 @@ pub trait StealthAddressOnCurve {
         expected_view_tag: u64,
     ) -> Option<Self::Fr> {
         let q_receiver = Self::compute_shared_point(viewing_key, ephemeral_public_key);
-
         let q_receiver_hashed = Self::hash_to_fr(&q_receiver.to_bytes());
-
-        // Check if retrieved view tag matches the expected view tag
-        let view_tag = q_receiver_hashed.get_view_tag();
-        if view_tag == expected_view_tag {
-            let stealth_private_key = spending_key + q_receiver_hashed;
-            Some(stealth_private_key)
+        if q_receiver_hashed.get_view_tag() == expected_view_tag {
+            Some(spending_key + q_receiver_hashed)
         } else {
             None
         }

sdk/src/vesta_impl.rs

@@ -1,5 +1,5 @@
-use crate::define_curve_tests;
-use crate::stealth_addresses::StealthAddressOnCurve;
+use crate::{define_curve_tests, stealth_addresses::StealthAddressOnCurve};
+/// Implementation of the StealthAddressOnCurve trait for the Vesta curve.
 use ark_vesta::{Fr, Projective};
 
 pub struct Vesta;