From b97723918051dc6c9b49a4931cdca047f946a23b Mon Sep 17 00:00:00 2001
From: 0xKanekiKen <100861945+0xKanekiKen@users.noreply.github.com>
Date: Mon, 23 Jan 2023 16:56:33 +0000
Subject: [PATCH] test: unit tests for merge in domain method

Signed-off-by: 0xKanekiKen <100861945+0xKanekiKen@users.noreply.github.com>
---
 src/hash/rpo/tests.rs | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/hash/rpo/tests.rs b/src/hash/rpo/tests.rs
index 03f69f94..80227ab1 100644
--- a/src/hash/rpo/tests.rs
+++ b/src/hash/rpo/tests.rs
@@ -1,5 +1,6 @@
 use super::{
-    Felt, FieldElement, Hasher, Rpo256, RpoDigest, StarkField, ALPHA, INV_ALPHA, STATE_WIDTH, ZERO,
+    Felt, FieldElement, Hasher, Rpo256, RpoDigest, StarkField, ALPHA, CAPACITY_RANGE, DIGEST_RANGE,
+    INV_ALPHA, RATE_RANGE, STATE_WIDTH, ZERO,
 };
 use core::convert::TryInto;
 use rand_utils::rand_value;
@@ -51,6 +52,40 @@ fn hash_elements_vs_merge() {
     assert_eq!(m_result, h_result);
 }
 
+#[test]
+fn hash_elements_vs_merge_in_domain() {
+    let elements = [Felt::new(rand_value()); 8];
+
+    let digests: [RpoDigest; 2] = [
+        RpoDigest::new(elements[..4].try_into().unwrap()),
+        RpoDigest::new(elements[4..].try_into().unwrap()),
+    ];
+
+    // pick a random domain value.
+    let domain = Felt::new(rand_value());
+
+    // convert the elements into a list of base field elements
+    let elements = Felt::as_base_elements(&elements);
+
+    // initialize state to all zeros.
+    let mut state = [ZERO; STATE_WIDTH];
+
+    // set the second capacity element to the domain.
+    state[CAPACITY_RANGE.start + 1] = domain;
+
+    // absorb elements into the state.
+    state[RATE_RANGE.start..RATE_RANGE.end].copy_from_slice(elements);
+
+    // apply permutation to the state.
+    Rpo256::apply_permutation(&mut state);
+
+    // return the first 4 elements of the state as hash result
+    let h_result = RpoDigest::new(state[DIGEST_RANGE].try_into().unwrap());
+
+    let m_result = Rpo256::merge_in_domain(&digests, domain);
+    assert_eq!(m_result, h_result);
+}
+
 #[test]
 fn hash_elements_vs_merge_with_int() {
     let tmp = [Felt::new(rand_value()); 4];