From 230311233b7d94bc5b56c0d3784d27054aa3bcef Mon Sep 17 00:00:00 2001
From: Georgi Grigorov <georgi.grigorov@limechain.tech>
Date: Tue, 29 Oct 2024 13:39:35 +0200
Subject: [PATCH] feat: adjusted claimPrimarySlot method in Authorship class

---
 .../java/com/limechain/babe/Authorship.java   | 28 +++++++++++--------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/main/java/com/limechain/babe/Authorship.java b/src/main/java/com/limechain/babe/Authorship.java
index 8c60f4ad..3c8409f1 100644
--- a/src/main/java/com/limechain/babe/Authorship.java
+++ b/src/main/java/com/limechain/babe/Authorship.java
@@ -1,5 +1,7 @@
 package com.limechain.babe;
 
+import com.limechain.babe.predigest.BabePreDigest;
+import com.limechain.babe.predigest.PreDigestType;
 import com.limechain.utils.LittleEndianUtils;
 import com.limechain.utils.math.BigRational;
 import com.limechain.chain.lightsyncstate.Authority;
@@ -17,15 +19,14 @@
 @NoArgsConstructor(access = AccessLevel.PRIVATE)
 public class Authorship {
 
-    //TODO: Replace return type with PreDigest
-    public static Object claimPrimarySlot(final byte[] randomness,
-                                          final long slotNumber,
-                                          final long epochNumber,
-                                          final Schnorrkel.KeyPair keyPair,
-                                          final int authorityIndex,
-                                          final BigInteger threshold) {
+    public static BabePreDigest claimPrimarySlot(final byte[] randomness,
+                                                 final BigInteger slotNumber,
+                                                 final long epochNumber,
+                                                 final Schnorrkel.KeyPair keyPair,
+                                                 final int authorityIndex,
+                                                 final BigInteger threshold) {
 
-        var transcript = makeTranscript(randomness, slotNumber, epochNumber);
+        var transcript = makeTranscript(randomness, slotNumber.longValue(), epochNumber);
 
         Schnorrkel schnorrkel = Schnorrkel.getInstance();
         VrfOutputAndProof vrfOutputAndProof = schnorrkel.vrfSign(keyPair, transcript);
@@ -38,8 +39,13 @@ public static Object claimPrimarySlot(final byte[] randomness,
         var isBelowThreshold = LittleEndianUtils.fromLittleEndianByteArray(vrfBytes).compareTo(threshold) < 0;
 
         if (isBelowThreshold) {
-            //TODO: Return PreDigest
-            return null;
+            return new BabePreDigest(
+                    PreDigestType.BABE_PRIMARY,
+                    authorityIndex,
+                    slotNumber,
+                    vrfOutputAndProof.getOutput(),
+                    vrfOutputAndProof.getProof()
+            );
         }
 
         return null;
@@ -99,7 +105,7 @@ private static double getBabeConstant(@NotNull Pair<BigInteger, BigInteger> cons
         return c;
     }
 
-    private static TranscriptData makeTranscript(byte[] randomness, Long slotNumber, Long epochNumber) {
+    private static TranscriptData makeTranscript(byte[] randomness, long slotNumber, long epochNumber) {
         var transcript = new TranscriptData("BABE".getBytes());
         transcript.appendMessage("slot number", LittleEndianUtils.longToLittleEndianBytes(slotNumber));
         transcript.appendMessage("current epoch", LittleEndianUtils.longToLittleEndianBytes(epochNumber));