⚡ Optimize find program key by adding static buffers and reusing the same SHA256 while finding an off-curve key

Author: Kuldotha

src/Solana.Unity.Dex/Orca/Address/AddressUtils.cs

@@ -20,28 +20,8 @@ public static class AddressUtils
         /// <returns>A Pda (program-derived address) off the curve, or null.</returns>
         public static Pda FindProgramAddress(IEnumerable<byte[]> seeds, PublicKey programId)
         {
-            byte nonce = 255;
-
-            while (nonce != 0)
-            {
-                PublicKey address;
-                List<byte[]> seedsWithNonce = new List<byte[]>(seeds);
-                seedsWithNonce.Add(new byte[] { nonce });
-
-                //try to generate the address 
-                bool created = PublicKey.TryCreateProgramAddress(new List<byte[]>(seedsWithNonce), programId, out address);
-
-                //if succeeded, return 
-                if (created)
-                {
-                    return new Pda(address, nonce);
-                }
-                
-                //decrease the nonce and retry if failed 
-                nonce--;
-            }
-
-            return null;
+            return PublicKey.TryFindProgramAddress(seeds, programId, out PublicKey pubkey, out byte bump) 
+                ? new Pda(pubkey, bump) : null;
         }
 
         /// <summary> 

src/Solana.Unity.Wallet/PublicKey.cs

@@ -256,44 +256,49 @@ public static bool IsValid(ReadOnlySpan<byte> key, bool validateCurve = false)
         /// <summary>
         /// The bytes of the `ProgramDerivedAddress` string.
         /// </summary>
-        private static readonly byte[] ProgramDerivedAddressBytes = Encoding.UTF8.GetBytes("ProgramDerivedAddress");
+        private static readonly byte[] ProgramDerivedAddressBytes = "ProgramDerivedAddress"u8.ToArray();
 
         /// <summary>
         /// Derives a program address.
         /// </summary>
         /// <param name="seeds">The address seeds.</param>
         /// <param name="programId">The program Id.</param>
         /// <param name="publicKey">The derived public key, returned as inline out.</param>
-        /// <returns>true if it could derive the program address for the given seeds, otherwise false..</returns>
+        /// <returns>true if it could derive the program address for the given seeds, otherwise false.</returns>
         /// <exception cref="ArgumentException">Throws exception when one of the seeds has an invalid length.</exception>
         public static bool TryCreateProgramAddress(ICollection<byte[]> seeds, PublicKey programId, out PublicKey publicKey)
         {
-            MemoryStream buffer = new(PublicKeyLength * seeds.Count + ProgramDerivedAddressBytes.Length + programId.KeyBytes.Length);
-
+            using SHA256 sha256 = SHA256.Create();
+            sha256.Initialize();
+            
             foreach (byte[] seed in seeds)
             {
                 if (seed.Length > PublicKeyLength)
                 {
                     throw new ArgumentException("max seed length exceeded", nameof(seeds));
                 }
-                buffer.Write(seed,0, seed.Length);
+                
+                sha256.TransformBlock(seed,0, seed.Length, null, 0);
             }
 
-            buffer.Write(programId.KeyBytes, 0, programId.KeyBytes.Length);
-            buffer.Write(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length);
-
-            SHA256 sha256 = SHA256.Create();
-            byte[] hash = sha256.ComputeHash(new ReadOnlySpan<byte>(buffer.GetBuffer(), 0, (int)buffer.Length).ToArray());
-
+            sha256.TransformBlock(programId.KeyBytes,0, programId.KeyBytes.Length,  null,0);
+            sha256.TransformBlock(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length, null, 0);
+            sha256.TransformFinalBlock([], 0, 0);
+            
+            byte[] hash = sha256.Hash!;
             if (hash.IsOnCurve())
             {
                 publicKey = null;
                 return false;
             }
-            publicKey = new(hash);
+
+            publicKey = new PublicKey(hash);
             return true;
         }
 
+        private static readonly List<byte[]> PdaSeedsBuffer = [];
+        private static readonly byte[] BumpArray = new byte[1];
+
         /// <summary>
         /// Attempts to find a program address for the passed seeds and program Id.
         /// </summary>
@@ -304,28 +309,42 @@ public static bool TryCreateProgramAddress(ICollection<byte[]> seeds, PublicKey
         /// <returns>True whenever the address for a nonce was found, otherwise false.</returns>
         public static bool TryFindProgramAddress(IEnumerable<byte[]> seeds, PublicKey programId, out PublicKey address, out byte bump)
         {
-            byte seedBump = 255;
-            List<byte[]> buffer = seeds.ToList();
-            var bumpArray = new byte[1];
-            buffer.Add(bumpArray);
+            PdaSeedsBuffer.Clear();
+            PdaSeedsBuffer.AddRange(seeds);
 
-            while (seedBump != 0)
+            if (PdaSeedsBuffer.Any(seed => seed.Length > PublicKeyLength))
+            {
+                throw new ArgumentException("max seed length exceeded", nameof(seeds));
+            }
+            
+            using SHA256 sha256 = SHA256.Create();
+            for (bump = 255; ; bump--)
             {
-                bumpArray[0] = seedBump;
-                bool success = TryCreateProgramAddress(buffer, programId, out PublicKey derivedAddress);
+                sha256.Initialize();
 
-                if (success)
+                foreach (byte[] seed in PdaSeedsBuffer)
+                {
+                    sha256.TransformBlock(seed, 0, seed.Length, null, 0);
+                }
+                
+                BumpArray[0] = bump;
+                sha256.TransformBlock(BumpArray, 0, 1, null, 0);
+                sha256.TransformBlock(programId.KeyBytes, 0, programId.KeyBytes.Length, null, 0);
+                sha256.TransformBlock(ProgramDerivedAddressBytes, 0, ProgramDerivedAddressBytes.Length, null, 0);
+                sha256.TransformFinalBlock([], 0, 0);
+                
+                byte[] hash = sha256.Hash!;
+                if (!hash.IsOnCurve())
                 {
-                    address = derivedAddress;
-                    bump = seedBump;
+                    address = new PublicKey(hash);
                     return true;
                 }
 
-                seedBump--;
+                if (bump == 0) 
+                    break;
             }
 
-            address = null;
-            bump = 0;
+            address = null!;
             return false;
         }