-
Notifications
You must be signed in to change notification settings - Fork 44
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Tree Optimization #120
Tree Optimization #120
Conversation
@@ -6,10 +6,14 @@ require ( | |||
github.com/golang/snappy v0.0.3 | |||
github.com/minio/sha256-simd v1.0.0 | |||
github.com/mitchellh/mapstructure v1.3.2 | |||
github.com/prysmaticlabs/gohashtree v0.0.1-alpha.0.20220714111606-acbb2962fb48 | |||
github.com/stretchr/testify v1.8.1 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I added stretchr
for testing assertions - it makes checking errors and equality a little simpler.
@@ -138,7 +138,7 @@ func (v *Value) hashTreeRoot(name string, appendBytes bool) string { | |||
}) | |||
} else { | |||
// dynamic bytes require special handling, need length mixed in | |||
hMethod := "PutBytes" | |||
hMethod := "Append" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The hash tree root of PreviousEpochParticipation
and CurrentEpochParticipation
does not work. Both PutBytes
and MerkleizeWithMixin
merkleizes the bytes if the bytes length is > 32, resulting in a double merkleization. The beacon spec tests doesn't uncover this bug because all the bytes in the spec tests for these two values are less than 32 bytes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can this happen in production? It might be because Prysm
does not use fastssz
for the BeaconState
and this case is not covered.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, it can happen in production. I believe that is correct. HashTreeRoot
for the BeaconState
crashed before it could get to this part of the code (it crashed at index 11, and this happens on index 15), so I think that's why we haven't encountered it before.
hasher.go
Outdated
if rest := len(input) % 32; rest != 0 { | ||
// pad zero bytes to the left | ||
input = append(input, zeroBytes[:32-rest]...) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Related to comment: https://github.com/ferranbt/fastssz/pull/120/files#r1097352609
PutBytes
and MerkleizeWithMixin
cannot be used together because they both merkleize bytes, resulting in double merkleization. The correct use is Append
with MerkleizeWithMixin
. However, MerkleizeWithMixin
should then check if the bytes should be padded to the next 32 bytes increment.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you use FillUpTo32
instead at the beginning of the function for the whole h.buf
array?
h.FillUpTo32()
input := h.buf[indx:]
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point! Addressed in b74a6c6.
// if byte list is empty, fill with zeros | ||
if len(b) == 0 { | ||
b = append(b, zeroBytes[:32]...) | ||
} | ||
// if byte list isn't filled with 32-bytes padded, pad | ||
if rest := len(b) % 32; rest != 0 { | ||
b = append(b, zeroBytes[:32-rest]...) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Similarly to https://github.com/ferranbt/fastssz/pull/120/files#r1097354675, we need to check if zero bytes need to added as padding to the next 32 byte increment.
func (w *Wrapper) fillEmptyNodes(i int) { | ||
func (w *Wrapper) getLimit(i int) int { | ||
size := len(w.nodes[i:]) | ||
for i := size; i < int(nextPowerOfTwo(uint64(size))); i++ { | ||
w.nodes = append(w.nodes, EmptyLeaf()) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We don't need to fill nodes anymore.
func TestBeaconStateTree_HashTreeRoot(t *testing.T) { | ||
data, err := os.ReadFile(TestFileName) | ||
require.NoError(t, err) | ||
|
||
sszState := BeaconStateBellatrix{} | ||
err = sszState.UnmarshalSSZ(data) | ||
require.NoError(t, err) | ||
|
||
tree, err := sszState.GetTree() | ||
require.NoError(t, err) | ||
|
||
hash := tree.Hash() | ||
|
||
require.Equal(t, "c4a9c5ebf637c089db599574b568bb679b385c1984f08410707db08e03d7ae52", hex.EncodeToString(hash)) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These tests are a sanity check against real data, as some of the spec tests didn't catch some issues with the generated encodings.
if structName == "BeaconState" || structName == "BeaconBlockBody" || structName == "ExecutionPayload" { | ||
// this gets to expensive, BeaconState even crashes with out-of-bounds memory allocation | ||
return | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This optimization allows running all the beacon spec tests.
// node with empty right node, add zero order hash as right node and mark right node as empty | ||
if nodes[leftIndex] != nil && nodes[rightIndex] == nil { | ||
nodes[i] = NewNodeWithLR(nodes[leftIndex], NewEmptyNode(zeroOrderHashes[k+1])) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Basically the crux of the optimization. 😄
leavesStart := powerTwo(depth) | ||
leafIndex := numLeaves - 1 | ||
|
||
nodes := make(map[int]*Node) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changes nodes
to a map instead of a slice so that we can lookup child nodes using the computed left and right index.
right *Node | ||
left *Node | ||
right *Node | ||
isEmpty bool |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Adds a flag so we know the zero order hash is actually an empty node and we shouldn't traverse the tree further for the right node.
hh.Append(b.PreviousEpochParticipation) | ||
hh.MerkleizeWithMixin(elemIndx, byteLen, (1099511627776+31)/32) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Prevents double merkleization of PutBytes
and MerkleizeWithMixin
used together.
This looks amazing! I will need a bit of time but It will most likely land by the end of the week. Could you regenerate the |
Thanks @ferranbt, that sounds good. CI should pass now, I tested on our fork: https://github.com/Snowfork/fastssz/actions/runs/4106825072/jobs/7085540042 I had to add |
@@ -5,7 +5,7 @@ make generate-testcases | |||
|
|||
# check differences | |||
cd sszgen/testcases | |||
if [[ `git status --porcelain` ]]; then | |||
if [[ `git status --porcelain .` ]]; then |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only check the current (sszgen/testcases
) directory, otherwise it picks up uncommitted changes in the rest of the project.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Approved! Some minor suggestions.
Co-authored-by: Ferran Borreguero <ferran.borreguero@gmail.com>
The current implementation of the FastSSZ tree fills empty nodes of the tree so that the number of tree leaves are to the power of 2. However, if the max capacity of a the tree leaves is very large (like BeaconState Validators, max capacity is
1 099 511 627 776
), Go cannot preallocate such a large slice (see issue #119 (comment)). Even if it could allocate the slice, adding all the empty nodes isn't necessary and slows down the tree creation since it needs to populate all the empty leaves, plus their parents.This PR changes the tree implementation to only add nodes that have values, and adds precomputed zero order hashes where the adjacent sibling requires a right-node sibling in order to hash the two nodes together.
Taking @protolambda's SSZ diagram as an example, in the current version of this library, nodes 7, 14 and 15 would be added as part of the tree. In this PR, only node 7 is added.
This improves the performance of
GetTree
significantly and enables the successful testing of all the spectests, whereBeaconBlock
,BeaconState
andExecutionPayload
were previously skipped because of performance issues.