From dbb32ff703de0dfe101ed018080af0c6a5b25f05 Mon Sep 17 00:00:00 2001
From: Song Xuyang <xuyangsong@hotmail.com>
Date: Mon, 4 Dec 2023 16:49:47 +0100
Subject: [PATCH] constrain the derivation of psi and rcm

---
 taiga_halo2/src/circuit/integrity.rs | 51 +++++++++++++++++++++++-----
 taiga_halo2/src/constant.rs          |  5 +++
 taiga_halo2/src/resource.rs          | 19 ++++++++---
 3 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/taiga_halo2/src/circuit/integrity.rs b/taiga_halo2/src/circuit/integrity.rs
index 3aa1a9b6..075d4fe7 100644
--- a/taiga_halo2/src/circuit/integrity.rs
+++ b/taiga_halo2/src/circuit/integrity.rs
@@ -4,7 +4,10 @@ use crate::circuit::{
     resource_commitment::{resource_commit, ResourceCommitChip},
     vp_circuit::{InputResourceVariables, OutputResourceVariables, ResourceVariables},
 };
-use crate::constant::{TaigaFixedBases, TaigaFixedBasesFull, POSEIDON_TO_CURVE_INPUT_LEN};
+use crate::constant::{
+    TaigaFixedBases, TaigaFixedBasesFull, POSEIDON_TO_CURVE_INPUT_LEN,
+    PRF_EXPAND_PERSONALIZATION_TO_FIELD, PRF_EXPAND_PSI, PRF_EXPAND_RCM,
+};
 use crate::resource::Resource;
 use crate::utils::poseidon_to_curve;
 use halo2_gadgets::{
@@ -233,19 +236,49 @@ pub fn check_output_resource(
         Value::known(output_resource.rseed),
     )?;
 
-    // TODO: constrain on psi and rcm derivation
     // Witness rcm
-    let rcm = assign_free_advice(
-        layouter.namespace(|| "witness rcm"),
+    let prf_expand_personalization = assign_free_constant(
+        layouter.namespace(|| "constant PRF_EXPAND_PERSONALIZATION_TO_FIELD"),
         advices[0],
-        Value::known(output_resource.get_rcm()),
+        *PRF_EXPAND_PERSONALIZATION_TO_FIELD,
+    )?;
+    let rcm_message = {
+        let prf_expand_rcm = assign_free_constant(
+            layouter.namespace(|| "constant PRF_EXPAND_RCM"),
+            advices[0],
+            pallas::Base::from(PRF_EXPAND_RCM as u64),
+        )?;
+        [
+            prf_expand_personalization.clone(),
+            prf_expand_rcm,
+            rseed.clone(),
+            old_nf.clone(),
+        ]
+    };
+    let rcm = poseidon_hash_gadget(
+        resource_commit_chip.get_poseidon_config(),
+        layouter.namespace(|| "derive the rcm"),
+        rcm_message,
     )?;
 
     // Witness psi
-    let psi = assign_free_advice(
-        layouter.namespace(|| "witness psi_output"),
-        advices[0],
-        Value::known(output_resource.get_psi()),
+    let psi_message = {
+        let prf_expand_psi = assign_free_constant(
+            layouter.namespace(|| "constant PRF_EXPAND_PSI"),
+            advices[0],
+            pallas::Base::from(PRF_EXPAND_PSI as u64),
+        )?;
+        [
+            prf_expand_personalization,
+            prf_expand_psi,
+            rseed.clone(),
+            old_nf.clone(),
+        ]
+    };
+    let psi = poseidon_hash_gadget(
+        resource_commit_chip.get_poseidon_config(),
+        layouter.namespace(|| "derive the psi"),
+        psi_message,
     )?;
 
     // Witness is_ephemeral
diff --git a/taiga_halo2/src/constant.rs b/taiga_halo2/src/constant.rs
index eb74f732..193e3ebf 100644
--- a/taiga_halo2/src/constant.rs
+++ b/taiga_halo2/src/constant.rs
@@ -27,6 +27,11 @@ pub const TRANSACTION_BINDING_HASH_PERSONALIZATION: &[u8; 16] = b"TxBindingSigHa
 pub const VP_COMMITMENT_PERSONALIZATION: &[u8; 8] = b"VPCommit";
 
 pub const PRF_EXPAND_PERSONALIZATION: &[u8; 16] = b"Taiga_ExpandSeed";
+lazy_static! {
+    pub static ref PRF_EXPAND_PERSONALIZATION_TO_FIELD: pallas::Base =
+        to_field_elements(PRF_EXPAND_PERSONALIZATION)[0];
+}
+
 pub const PRF_EXPAND_PSI: u8 = 0;
 pub const PRF_EXPAND_RCM: u8 = 1;
 pub const PRF_EXPAND_PUBLIC_INPUT_PADDING: u8 = 2;
diff --git a/taiga_halo2/src/resource.rs b/taiga_halo2/src/resource.rs
index 84041199..8c67dd27 100644
--- a/taiga_halo2/src/resource.rs
+++ b/taiga_halo2/src/resource.rs
@@ -4,8 +4,9 @@ use crate::{
         vp_examples::{TrivialValidityPredicateCircuit, COMPRESSED_TRIVIAL_VP_VK},
     },
     constant::{
-        NUM_RESOURCE, POSEIDON_TO_CURVE_INPUT_LEN, PRF_EXPAND_PERSONALIZATION, PRF_EXPAND_PSI,
-        PRF_EXPAND_PUBLIC_INPUT_PADDING, PRF_EXPAND_RCM, PRF_EXPAND_VCM_R,
+        NUM_RESOURCE, POSEIDON_TO_CURVE_INPUT_LEN, PRF_EXPAND_PERSONALIZATION,
+        PRF_EXPAND_PERSONALIZATION_TO_FIELD, PRF_EXPAND_PSI, PRF_EXPAND_PUBLIC_INPUT_PADDING,
+        PRF_EXPAND_RCM, PRF_EXPAND_VCM_R,
     },
     merkle_tree::{Anchor, MerklePath, Node},
     nullifier::{Nullifier, NullifierKeyContainer},
@@ -270,12 +271,22 @@ impl Resource {
 
     // psi is the randomness used to derive the nullifier
     pub fn get_psi(&self) -> pallas::Base {
-        poseidon_hash_n([self.rseed, self.nonce.inner()])
+        poseidon_hash_n([
+            *PRF_EXPAND_PERSONALIZATION_TO_FIELD,
+            pallas::Base::from(PRF_EXPAND_PSI as u64),
+            self.rseed,
+            self.nonce.inner(),
+        ])
     }
 
     // rcm is the randomness of resource commitment
     pub fn get_rcm(&self) -> pallas::Base {
-        poseidon_hash_n([self.rseed, self.nonce.inner()])
+        poseidon_hash_n([
+            *PRF_EXPAND_PERSONALIZATION_TO_FIELD,
+            pallas::Base::from(PRF_EXPAND_RCM as u64),
+            self.rseed,
+            self.nonce.inner(),
+        ])
     }
 
     pub fn calculate_root(&self, path: &MerklePath) -> Anchor {