From 88462cf654583613166c295537dfeed55219e57d Mon Sep 17 00:00:00 2001
From: Alvaro Revuelta <alvrevuelta@gmail.com>
Date: Fri, 31 May 2024 07:11:54 +0200
Subject: [PATCH] Add InsertRawLeaf (#23)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: richΛrd <info@richardramos.me>
---
 rln/rln.go      |  8 ++++++++
 rln/rln_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/rln/rln.go b/rln/rln.go
index 7049c54..911512a 100644
--- a/rln/rln.go
+++ b/rln/rln.go
@@ -427,6 +427,14 @@ func (r *RLN) InsertMember(idComm IDCommitment, userMessageLimit uint32) error {
 	return nil
 }
 
+func (r *RLN) InsertRawLeaf(rawLeaf MerkleNode) error {
+	insertionSuccess := r.w.SetNextLeaf(rawLeaf[:])
+	if !insertionSuccess {
+		return errors.New("could not insert raw leaf")
+	}
+	return nil
+}
+
 // Insert multiple members i.e., identity commitments starting from index
 // This proc is atomic, i.e., if any of the insertions fails, all the previous insertions are rolled back
 func (r *RLN) InsertMembers(index MembershipIndex, idComms []IDCommitment) error {
diff --git a/rln/rln_test.go b/rln/rln_test.go
index 6fbace9..a1da4c1 100644
--- a/rln/rln_test.go
+++ b/rln/rln_test.go
@@ -80,6 +80,37 @@ func (s *RLNSuite) TestInsertMember() {
 	s.NoError(err)
 }
 
+func (s *RLNSuite) TestInsertRawLeaf() {
+	rln, err := NewRLN()
+	s.NoError(err)
+
+	for i := 0; i < 10; i++ {
+		// Generate a membership
+		memKeys, err := rln.MembershipKeyGen(10)
+		s.NoError(err)
+
+		// Calculate the leaf ourselves
+		userMessageLimitBytes := SerializeUint32(memKeys.UserMessageLimit)
+		hashedLeaf, err := rln.Poseidon(memKeys.IDCommitment[:], userMessageLimitBytes[:])
+		s.NoError(err)
+
+		// Insert the leaf as it is
+		err = rln.InsertRawLeaf(hashedLeaf)
+		s.NoError(err)
+
+		// Get it from the tree
+		retrievedLeaf, err := rln.GetLeaf(uint(i))
+		s.NoError(err)
+
+		// Check the retrieved matches the one we added
+		s.Equal(hashedLeaf, retrievedLeaf)
+
+		// Check tree size matches
+		numLeaves := rln.LeavesSet()
+		s.Equal(uint(i+1), numLeaves)
+	}
+}
+
 func (s *RLNSuite) TestInsertMembers() {
 	rln, err := NewRLN()
 	s.NoError(err)