Dev/jgough/10207 support extrabytes (#37)

* Update AddHashedLeaf to include extra bytes for trie value

re: AB#10207

---------

Co-authored-by: jgough <[email protected]>

massifs/massifcontext.go

@@ -369,10 +369,15 @@ func (mc *MassifContext) Append(value []byte) (uint64, error) {
 // trie. On error, the current data buffer should be discarded entirely (not
 // written back to storage)
 //
+// Params:
+//   - extraBytes - extra bytes that are added to the trie value before idtimestamp. maximum 24 bytes.
+//     any extra bytes above 24 bytes will be truncated.
+//
 // Returns the resulting size of the mmr if the leaf is adds successfully.
 func (mc *MassifContext) AddHashedLeaf(
 	hasher hash.Hash,
 	idTimestamp uint64,
+	extraBytes []byte,
 	logId []byte,
 	appId []byte,
 	value []byte,
@@ -406,7 +411,7 @@ func (mc *MassifContext) AddHashedLeaf(
 	nextLeafIndex := mc.MassifLeafCount()
 
 	// Overwrite the pre-allocated index entry with the index data.
-	SetTrieEntry(mc.Data, mc.IndexStart(), nextLeafIndex, idTimestamp, trieKey)
+	SetTrieEntry(mc.Data, mc.IndexStart(), nextLeafIndex, idTimestamp, extraBytes, trieKey)
 
 	// Save the last id added so that we can guarantee monotonicity (and hence uniqueness for the tenant)
 	mc.setLastIdTimestamp(idTimestamp)

massifs/testcommitter.go

@@ -162,7 +162,7 @@ func (c *TestMinimalCommitter) AddLeaves(
 
 	for _, args := range batch {
 
-		_, err = mc.AddHashedLeaf(hasher, args.Id, args.LogId, args.AppId, args.Value)
+		_, err = mc.AddHashedLeaf(hasher, args.Id, args.LogId, args.AppId, nil, args.Value)
 		if errors.Is(err, ErrMassifFull) {
 			_, err = c.committer.CommitContext(ctx, mc)
 			if err != nil {
@@ -181,7 +181,7 @@ func (c *TestMinimalCommitter) AddLeaves(
 
 			// Remember to add the leaf we failed to add above
 			_, err = mc.AddHashedLeaf(
-				hasher, args.Id, args.LogId, args.AppId, args.Value)
+				hasher, args.Id, args.LogId, args.AppId, nil, args.Value)
 			if err != nil {
 				c.log.Infof("AddLeaves: %v", err)
 				return err

massifs/trieentry.go

@@ -23,12 +23,28 @@ import (
 )
 
 const (
-	TrieEntryBytes            = 32 * 2
+
+	/**
+	 * Each Trie Entry is the following:
+	 *
+	 * |----------|-------------|--------------|
+	 * | Trie Key | Extra Bytes | ID Timestamp |
+	 * |----------|-------------|--------------|
+	 * | 32 bytes |  24 bytes   |    8 bytes   |
+	 * |----------|-------------|--------------|
+	 *
+	 * Where Trie Value = Extra Bytes + ID Timestamp
+	 */
+
+	TrieEntryBytes            = 32 * 2 // 32 for trie key and 32 for trie value
 	TrieKeyBytes              = 32
 	TrieKeyEnd                = TrieKeyBytes
 	TrieEntryIdTimestampStart = 32 + 24
 	TrieEntrySnowflakeIDBytes = 8
 	TrieEntryIdTimestampEnd   = TrieEntryIdTimestampStart + TrieEntrySnowflakeIDBytes
+	TrieEntryExtraBytesStart  = 32
+	TrieEntryExtraBytesSize   = 24
+	TrieEntryExtraBytesEnd    = TrieEntryExtraBytesStart + TrieEntryExtraBytesSize
 )
 
 var (
@@ -129,6 +145,20 @@ func GetIdtimestamp(trieData []byte, indexStart uint64, trieIndex uint64) []byte
 	return trieData[idStart:idEnd]
 }
 
+// GetExtraBytes returns the extra bytes corresponding to the given trie index,
+// from the given trie data.
+//
+// extraBytes are part of the trie value, where trieValue = extraBytes + idtimestamp
+//
+// NOTE: trieIndex is equivilent to leafIndex.
+func GetExtraBytes(trieData []byte, indexStart uint64, trieIndex uint64) []byte {
+	trieEntryOffset := TrieEntryOffset(indexStart, trieIndex)
+	extraBytesStart := trieEntryOffset + TrieEntryExtraBytesStart
+	extraBytesEnd := trieEntryOffset + TrieEntryExtraBytesEnd
+
+	return trieData[extraBytesStart:extraBytesEnd]
+}
+
 // SetTrieEntry stores the trie Entry (trieKey + idTimestamp) in the given trie data at the given
 //
 //	trie index.
@@ -137,14 +167,22 @@ func GetIdtimestamp(trieData []byte, indexStart uint64, trieIndex uint64) []byte
 //
 //	for leaves.
 func SetTrieEntry(trieData []byte, indexStart uint64, trieIndex uint64,
-	idTimestamp uint64, trieKey []byte) {
+	idTimestamp uint64, extraBytes []byte, trieKey []byte) {
 
 	trieEntryOffset := TrieEntryOffset(indexStart, trieIndex)
 	copy(trieData[trieEntryOffset:trieEntryOffset+TrieKeyEnd], trieKey)
 
-	idStart := trieEntryOffset + TrieEntryIdTimestampStart
-	idEnd := trieEntryOffset + TrieEntryIdTimestampEnd
-	binary.BigEndian.PutUint64(trieData[idStart:idEnd], idTimestamp)
+	// extra bytes
+	if extraBytes != nil {
+		extraBytesStart := trieEntryOffset + TrieEntryExtraBytesStart
+		extraBytesEnd := trieEntryOffset + TrieEntryExtraBytesEnd
+		copy(trieData[extraBytesStart:extraBytesEnd], extraBytes)
+	}
+
+	// idtimestamp
+	idTimestampStart := trieEntryOffset + TrieEntryIdTimestampStart
+	idTimestampEnd := trieEntryOffset + TrieEntryIdTimestampEnd
+	binary.BigEndian.PutUint64(trieData[idTimestampStart:idTimestampEnd], idTimestamp)
 }
 
 // NewTrieKey creates the trie key value.

massifs/trieentry_test.go

@@ -123,6 +123,7 @@ func TestSetLogIndexEntry(t *testing.T) {
 		indexStart  uint64
 		leafIndex   uint64
 		idTimestamp uint64
+		extraBytes  []byte
 		index       []byte
 		before      []byte
 		after       []byte
@@ -138,6 +139,7 @@ func TestSetLogIndexEntry(t *testing.T) {
 				indexStart:  uint64(expectTrieEntryBytes * 2),
 				leafIndex:   1, // *trie* index NOT mmrIndex,
 				idTimestamp: 0x0102030405060708,
+				extraBytes:  []byte(`888888888888888888888888`), // maxiumum size of 24 bytes
 				index:       b64one[:32],
 				before:      b64zero[:32],
 				after:       b64two[:32],
@@ -150,6 +152,20 @@ func TestSetLogIndexEntry(t *testing.T) {
 				indexStart:  uint64(expectTrieEntryBytes * 2),
 				leafIndex:   3, // *trie* index NOT mmrIndex
 				idTimestamp: 0x0102030405060708,
+				extraBytes:  []byte(`8888888888888888888888889999`), // overflow maxiumum size of 24 bytes (should truncate the 9's)
+				index:       b64three[:32],
+				before:      b64two[:32],
+				after:       b64four[:32],
+			},
+		},
+		{
+			name: "get trie index 3, nil extra bytes",
+			args: args{
+				logData:     slices.Concat(b64h0, b64h1, b64zero, b64one, b64two, b64three, b64four),
+				indexStart:  uint64(expectTrieEntryBytes * 2),
+				leafIndex:   3, // *trie* index NOT mmrIndex
+				idTimestamp: 0x0102030405060708,
+				extraBytes:  nil,
 				index:       b64three[:32],
 				before:      b64two[:32],
 				after:       b64four[:32],
@@ -165,7 +181,7 @@ func TestSetLogIndexEntry(t *testing.T) {
 			if tt.args.before == nil && tt.args.after == nil {
 				return
 			}
-			SetTrieEntry(tt.args.logData, tt.args.indexStart, tt.args.leafIndex, tt.args.idTimestamp, b64clear)
+			SetTrieEntry(tt.args.logData, tt.args.indexStart, tt.args.leafIndex, tt.args.idTimestamp, tt.args.extraBytes, b64clear)
 
 			gotBefore := GetTrieKey(tt.args.logData, tt.args.indexStart, tt.args.leafIndex-1)
 			assert.Equal(t, tt.args.before, gotBefore)
@@ -179,6 +195,14 @@ func TestSetLogIndexEntry(t *testing.T) {
 
 			gotId := GetIdtimestamp(tt.args.logData, tt.args.indexStart, tt.args.leafIndex)
 			assert.Equal(t, tt.args.idTimestamp, binary.BigEndian.Uint64(gotId))
+
+			// check the extra bytes
+			gotBytes := GetExtraBytes(tt.args.logData, tt.args.indexStart, tt.args.leafIndex)
+			assert.Equal(t, 24, len(gotBytes)) // make doubley sure we have filled the extra bytes and there is no overflow
+
+			if tt.args.extraBytes != nil {
+				assert.Equal(t, tt.args.extraBytes[0:24], gotBytes)
+			}
 		})
 	}
 }