Skip to content

Commit

Permalink
add nibbleset in RlpMemo to remember the order
Browse files Browse the repository at this point in the history
  • Loading branch information
@dipkakwani
dipkakwani committed Jan 6, 2025
1 parent 17d01e2 commit bd369fe
Show file tree
Hide file tree
Showing 5 changed files with 189 additions and 169 deletions.
114 changes: 68 additions & 46 deletions src/Paprika.Tests/Merkle/RlpMemoTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ public class RlpMemoTests
private enum RlpMemoOperation
{
Set,
SetRaw,
Clear,
Delete,
Insert
Expand All @@ -25,11 +24,17 @@ public void Random_delete()
Span<byte> raw = stackalloc byte[RlpMemo.MaxSize];
var children = new NibbleSet();

for (var i = 0; i < RlpMemo.MaxSize; i++)
for (var i = 0; i < RlpMemo.MaxSize - NibbleSet.MaxByteSize; i++)
{
raw[i] = (byte)(i & 0xFF);
}

// Set all the index bits at the end.
for (var i = RlpMemo.MaxSize - 1; i >= RlpMemo.MaxSize - NibbleSet.MaxByteSize; i--)
{
raw[i] = 0xFF;
}

for (var i = 0; i < NibbleSet.NibbleCount; i++)
{
children[(byte)i] = true;
Expand All @@ -48,11 +53,15 @@ public void Random_delete()
}

children[child] = false;
memo = RlpMemo.Delete(memo, child, children, raw);
memo = RlpMemo.Delete(memo, child, raw);

memo.Length.Should().Be(RlpMemo.MaxSize - (i + 1) * Keccak.Size);
memo.Exists(child, children).Should().BeFalse();
memo.TryGetKeccak(child, out var keccak, children).Should().BeFalse();
var expectedLength = (i != NibbleSet.NibbleCount - 1)
? RlpMemo.MaxSize - (i + 1) * Keccak.Size + NibbleSet.MaxByteSize
: 0;

memo.Length.Should().Be(expectedLength);
memo.Exists(child).Should().BeFalse();
memo.TryGetKeccak(child, out var keccak).Should().BeFalse();
keccak.IsEmpty.Should().BeTrue();
}

Expand Down Expand Up @@ -87,11 +96,15 @@ public void Random_insert()
}

children[child] = true;
memo = RlpMemo.Insert(memo, child, children, keccak, workingMemory);
memo = RlpMemo.Insert(memo, child, keccak, workingMemory);

var expectedLength = (i != NibbleSet.NibbleCount - 1)
? (i + 1) * Keccak.Size + NibbleSet.MaxByteSize
: RlpMemo.MaxSize;

memo.Length.Should().Be((i + 1) * Keccak.Size);
memo.Exists(child, children).Should().BeTrue();
memo.TryGetKeccak(child, out var k, children).Should().BeTrue();
memo.Length.Should().Be(expectedLength);
memo.Exists(child).Should().BeTrue();
memo.TryGetKeccak(child, out var k).Should().BeTrue();
k.SequenceEqual(keccak).Should().BeTrue();
}

Expand All @@ -102,29 +115,39 @@ public void Random_insert()
public void In_place_update()
{
Span<byte> raw = stackalloc byte[RlpMemo.MaxSize];
var children = new NibbleSet();

for (var i = 0; i < RlpMemo.MaxSize; i++)
for (var i = 0; i < RlpMemo.MaxSize - NibbleSet.MaxByteSize; i++)
{
raw[i] = (byte)(i & 0xFF);
}

var memo = new RlpMemo(raw);
// Set all the index bits at the end.
for (var i = RlpMemo.MaxSize - 1; i >= RlpMemo.MaxSize - NibbleSet.MaxByteSize; i--)
{
raw[i] = 0xFF;
}

var children = new NibbleSet();
for (var i = 0; i < NibbleSet.NibbleCount; i++)
{
children[(byte)i] = true;
}

var memo = new RlpMemo(raw);

// Delete each child and the corresponding keccak
for (byte i = 0; i < NibbleSet.NibbleCount; i++)
{
children[i] = false;
memo = RlpMemo.Delete(memo, i, children, raw);
memo = RlpMemo.Delete(memo, i, raw);

var expectedLength = (i != NibbleSet.NibbleCount - 1)
? RlpMemo.MaxSize - (i + 1) * Keccak.Size + NibbleSet.MaxByteSize
: 0;

memo.Length.Should().Be(RlpMemo.MaxSize - (i + 1) * Keccak.Size);
memo.Exists(i, children).Should().BeFalse();
memo.TryGetKeccak(i, out var k, children).Should().BeFalse();
memo.Length.Should().Be(expectedLength);
memo.Exists(i).Should().BeFalse();
memo.TryGetKeccak(i, out var k).Should().BeFalse();
k.IsEmpty.Should().BeTrue();
}

Expand All @@ -137,11 +160,15 @@ public void In_place_update()
for (byte i = 0; i < NibbleSet.NibbleCount; i++)
{
children[i] = true;
memo = RlpMemo.Insert(memo, i, children, keccak, raw);
memo = RlpMemo.Insert(memo, i, keccak, raw);

memo.Length.Should().Be((i + 1) * Keccak.Size);
memo.Exists(i, children).Should().BeTrue();
memo.TryGetKeccak(i, out var k, children).Should().BeTrue();
var expectedLength = (i != NibbleSet.NibbleCount - 1)
? (i + 1) * Keccak.Size + NibbleSet.MaxByteSize
: RlpMemo.MaxSize;

memo.Length.Should().Be(expectedLength);
memo.Exists(i).Should().BeTrue();
memo.TryGetKeccak(i, out var k).Should().BeTrue();
k.SequenceEqual(keccak).Should().BeTrue();
}

Expand Down Expand Up @@ -173,14 +200,19 @@ public void Large_random_operations(int numOperations)
var children = new NibbleSet();

Span<byte> keccak = new byte[Keccak.Size];
keccak.Fill(0xFF);
KeccakOrRlp.FromSpan(keccak, out var keccakOrRlp);
keccak.Fill(0xFF);

for (var i = 0; i < RlpMemo.MaxSize; i++)
for (var i = 0; i < RlpMemo.MaxSize - NibbleSet.MaxByteSize; i++)
{
raw[i] = (byte)(i & 0xFF);
}

// Set all the index bits at the end.
for (var i = RlpMemo.MaxSize - 1; i >= RlpMemo.MaxSize - NibbleSet.MaxByteSize; i--)
{
raw[i] = 0xFF;
}

for (var i = 0; i < NibbleSet.NibbleCount; i++)
{
children[(byte)i] = true;
Expand All @@ -197,53 +229,43 @@ public void Large_random_operations(int numOperations)
switch (op)
{
case RlpMemoOperation.Set:
if (memo.Exists(child, children))
{
memo.Set(keccakOrRlp, child, children);

memo.TryGetKeccak(child, out var k, children).Should().BeTrue();
k.SequenceEqual(keccakOrRlp.Span).Should().BeTrue();
}

break;
case RlpMemoOperation.SetRaw:
if (memo.Exists(child, children))
if (memo.Exists(child))
{
memo.SetRaw(keccak, child, children);
memo.Set(keccak, child);

memo.TryGetKeccak(child, out var k, children).Should().BeTrue();
memo.TryGetKeccak(child, out var k).Should().BeTrue();
k.SequenceEqual(keccak).Should().BeTrue();
}

break;
case RlpMemoOperation.Clear:
if (memo.Exists(child, children))
if (memo.Exists(child))
{
memo.Clear(child, children);
memo.Clear(child);

memo.TryGetKeccak(child, out var k, children).Should().BeFalse();
memo.TryGetKeccak(child, out var k).Should().BeFalse();
k.IsEmpty.Should().BeTrue();
}

break;
case RlpMemoOperation.Delete:
if (memo.Exists(child, children))
if (memo.Exists(child))
{
children[child] = false;
memo = RlpMemo.Delete(memo, child, children, raw);
memo = RlpMemo.Delete(memo, child, raw);

memo.TryGetKeccak(child, out var k, children).Should().BeFalse();
memo.TryGetKeccak(child, out var k).Should().BeFalse();
k.IsEmpty.Should().BeTrue();
}

break;
case RlpMemoOperation.Insert:
if (!memo.Exists(child, children))
if (!memo.Exists(child))
{
children[child] = true;
memo = RlpMemo.Insert(memo, child, children, keccak, raw);
memo = RlpMemo.Insert(memo, child, keccak, raw);

memo.TryGetKeccak(child, out var k, children).Should().BeTrue();
memo.TryGetKeccak(child, out var k).Should().BeTrue();
k.SequenceEqual(keccak).Should().BeTrue();
}

Expand Down
1 change: 0 additions & 1 deletion src/Paprika/Merkle/CommitExtensions.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@ public static void SetBranch(this ICommit commit, in Key key, NibbleSet.Readonly
EntryType type = EntryType.Persistent)
{
var branch = new Branch(children);
Debug.Assert(rlp.Length == 0 || rlp.Length == (children.SetCount * Keccak.Size));
commit.Set(key, branch.WriteTo(stackalloc byte[Branch.MaxByteLength]), rlp, type);
}

Expand Down
63 changes: 21 additions & 42 deletions src/Paprika/Merkle/ComputeMerkleBehavior.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -277,39 +277,19 @@ private BuildStorageTriesItem[] GetStorageWorkItems(ICommitWithStats commit, Cac

public ReadOnlySpan<byte> InspectBeforeApply(in Key key, ReadOnlySpan<byte> data, Span<byte> workingSet)
{
if (data.IsEmpty)
return data;

if (key.Type != DataType.Merkle)
return data;

var node = Node.Header.Peek(data).NodeType;

if (node != Node.Type.Branch)
if (data.IsEmpty || key.Type != DataType.Merkle)
{
// Return data as is, either the node is not a branch or the memoization is not set for branches.
return data;
}

if (key.Path.Length < SkipRlpMemoizationForTopLevelsCount)
if (key.Path.Length < SkipRlpMemoizationForTopLevelsCount &&
Node.Header.Peek(data).NodeType == Node.Type.Branch)
{
// For State branches, omit top levels of RLP memoization
return Node.Branch.GetOnlyBranchData(data);
}

var memoizedRlp = Node.Branch.ReadFrom(data, out var branch);
if (memoizedRlp.Length == 0)
{
// no RLP of children memoized, return
return data;
}

// Copy the memoized RLPs
var dataLength = data.Length - memoizedRlp.Length;
data[..dataLength].CopyTo(workingSet);
memoizedRlp.CopyTo(workingSet[dataLength..]);

return workingSet[..data.Length];
return data;
}

public Keccak RootHash { get; private set; }
Expand Down Expand Up @@ -511,7 +491,7 @@ private void EncodeBranch(scoped in Key key, scoped in ComputeContext ctx, scope
{
if (branch.Children[i])
{
if (memoize && memo.TryGetKeccak(i, out var keccak, branch.Children))
if (memoize && memo.TryGetKeccak(i, out var keccak))
{
// keccak from cache
stream.Encode(keccak);
Expand Down Expand Up @@ -544,23 +524,23 @@ private void EncodeBranch(scoped in Key key, scoped in ComputeContext ctx, scope
{
memoizedUpdated = true;

if (memo.Exists(i, branch.Children))
if (memo.Exists(i))
{
memo.Set(keccakOrRlp, i, branch.Children);
memo.Set(keccakOrRlp.Span, i);
}
else if (keccakOrRlp.DataType == KeccakOrRlp.Type.Keccak)
{
memo = RlpMemo.Insert(memo, i, branch.Children, keccakOrRlp.Span, rlpMemoization);
memo = RlpMemo.Insert(memo, i, keccakOrRlp.Span, rlpMemoization);
}
}
}
else
{
stream.EncodeEmptyArray();

if (memoize && memo.Exists(i, branch.Children))
if (memoize && memo.Exists(i))
{
memo = RlpMemo.Delete(memo, i, branch.Children, rlpMemoization);
memo = RlpMemo.Delete(memo, i, rlpMemoization);
memoizedUpdated = true;
}
}
Expand Down Expand Up @@ -608,18 +588,18 @@ private void EncodeBranch(scoped in Key key, scoped in ComputeContext ctx, scope

if (value.Length == Keccak.Size)
{
if (memo.Exists(i, branch.Children))
if (memo.Exists(i))
{
memo.SetRaw(value, i, branch.Children);
memo.Set(value, i);
}
else
{
memo = RlpMemo.Insert(memo, i, branch.Children, value, rlpMemoization);
memo = RlpMemo.Insert(memo, i, value, rlpMemoization);
}
}
else if (memo.Exists(i, branch.Children))
else if (memo.Exists(i))
{
memo.Clear(i, branch.Children);
memo = RlpMemo.Delete(memo, i, rlpMemoization);
}
}
}
Expand Down Expand Up @@ -834,7 +814,6 @@ private static DeleteStatus Delete(in NibblePath path, int at, ICommit commit, C
case Node.Type.Branch:
{
var nibble = path[at];
Debug.Assert(leftover.Length == 0 || leftover.Length == (branch.Children.SetCount * Keccak.Size));
if (!branch.Children[nibble])
{
// no such child
Expand Down Expand Up @@ -939,18 +918,18 @@ static void UpdateBranchOnDelete(ICommit commit, in Node.Branch branch, NibbleSe
}

// If this child still exists, only clear the memo. Otherwise, delete it from the memo.
if (memo.Exists(nibble, children))
if (children[nibble] && memo.Exists(nibble))
{
memo.Clear(nibble, children);
memo.Clear(nibble);
}
else if (leftover.Length > 0)
else if (memo.Exists(nibble))
{
if (rlpWorkingSet == null)
{
rlpWorkingSet = ArrayPool<byte>.Shared.Rent(leftover.Length - Keccak.Size);
}

memo = RlpMemo.Delete(memo, nibble, children, rlpWorkingSet);
memo = RlpMemo.Delete(memo, nibble, rlpWorkingSet);
}

var shouldUpdate = !branch.Children.Equals(children);
Expand Down Expand Up @@ -1218,11 +1197,11 @@ private static void MarkPathDirty(in NibblePath path, in Span<byte> rlpMemoWorki
// If this path does not exist, insert an empty Keccak. Otherwise, clear it in the memo.
if (createLeaf)
{
memo = RlpMemo.Insert(memo, nibble, children, Keccak.Zero.Span, rlpMemoWorkingSet);
memo = RlpMemo.Insert(memo, nibble, Keccak.Zero.Span, rlpMemoWorkingSet);
}
else
{
memo.Clear(nibble, children);
memo.Clear(nibble);
}

if (shouldUpdateBranch || childRlpRequiresUpdate)
Expand Down
Loading

0 comments on commit bd369fe

Please sign in to comment.