Fix encoding for branch and extension for inlined data.

src/Paprika.Tests/Merkle/Commit.cs

@@ -28,6 +28,9 @@ public class Commit(bool skipMemoizedRlpCheck = false) : ICommit
     public void Set(in Key key, ReadOnlySpan<byte> value)
     {
         _before[GetKey(key)] = value.ToArray();
+        //to enable storage root calculation for tests
+        if (!key.IsState)
+            _stats.RegisterSetStorageAccount(key.Path.UnsafeAsKeccak);
     }
 
     public void DeleteKey(in Key key) => Set(key, ReadOnlySpan<byte>.Empty);

src/Paprika.Tests/Merkle/RootHashTests.cs

@@ -88,6 +88,28 @@ public void Extension()
         AssertRoot("a624947d9693a5cba0701897b3a48cb9954c2f4fd54de36151800eb2c7f6bf50", commit);
     }
 
+    [Test]
+    public void Extension_branch_two_short_leaves()
+    {
+        var commit = new Commit();
+
+        //create E->B->L
+        //           ->L
+        //leaves without any key and very small value cause to be inlined in branch
+        //encoded branch rlp is also < 32 bytes which causes it to be encoded as RLP in extension node
+        Keccak storageKey1 =
+            new Keccak(Convert.FromHexString("ccccccccccccccccccccddddddddddddddddeeeeeeeeeeeeeeeeeeeeeeeeeeb1"));
+        Keccak storageKey2 =
+            new Keccak(Convert.FromHexString("ccccccccccccccccccccddddddddddddddddeeeeeeeeeeeeeeeeeeeeeeeeeeb2"));
+
+        commit.Set(Key.Account(Values.Key0),
+            new Account(0, 1).WriteTo(stackalloc byte[Paprika.Account.MaxByteCount]));
+        commit.Set(Key.StorageCell(NibblePath.FromKey(Values.Key0), storageKey1), new byte[] { 1, 2, 3 });
+        commit.Set(Key.StorageCell(NibblePath.FromKey(Values.Key0), storageKey2), new byte[] { 10, 20, 30 });
+
+        AssertRoot("c8f9c9c3a0e95e13d6f6b7e0df65052a0cc484ab0db3e57c287df74f2714d5b3", commit);
+    }
+
     [TestCase(0, "7f7fd47a28dc4dbfd1b1b33d254da8be74deab55bef81a02c232ca9957e05689", TestName = "From the Root")]
     [TestCase(Keccak.Size - 2, "d8fc42b5f9491f526d0935445e9b83d8ddde46978cc450a6d1f83351da1bfae2",
         TestName = "At the bottom")]

src/Paprika/Merkle/ComputeMerkleBehavior.cs

@@ -649,7 +649,9 @@ private void EncodeBranch(scoped in Key key, scoped in ComputeContext ctx, scope
         // Write length of length in front of the payload, resetting the stream properly
         var end = stream.Position;
         var actualLength = end - initialShift;
-        var lengthOfLength = Rlp.LengthOfLength(actualLength) + 1;
+        var lengthOfLength = Rlp.LengthOfLength(actualLength);
+        if (actualLength >= 56) //to match StartSequence
+            lengthOfLength++;
         var from = initialShift - lengthOfLength;
         stream.Position = from;
         stream.StartSequence(actualLength);
@@ -692,7 +694,10 @@ private void EncodeExtension(scoped in Key key, scoped in ComputeContext ctx, sc
         RlpStream stream = new(pooled.Span.Slice(slice, totalLength));
         stream.StartSequence(contentLength);
         stream.Encode(span);
-        stream.Encode(keccakOrRlp.Keccak);
+        if (keccakOrRlp.DataType == KeccakOrRlp.Type.Rlp)
+            stream.Write(keccakOrRlp.Span);
+        else
+            stream.Encode(keccakOrRlp.Keccak);
         stream.ToKeccakOrRlp(out keccakOrRlp);
     }