Skip to content

Commit b264ff6

Browse files
fxamackerfxamacker
authored
Merge pull request #629 from onflow/fxamacker/add-array-and-byte-slice-conversion
Add two new functions to optimize atree array and Go []byte conversions
2 parents d1d06a9 + ae51900 commit b264ff6

File tree

5 files changed

+545
-181
lines changed

5 files changed

+545
-181
lines changed

‎array_conversion.go‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
/*
2+
* Atree - Scalable Arrays and Ordered Maps
3+
*
4+
* Copyright Flow Foundation
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the "License");
7+
* you may not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package atree
20+
21+
import (
22+
"reflect"
23+
)
24+
25+
// ByteStorable is a type constraint that permits any type that
26+
// implements Storable interface and has byte underlying type.
27+
type ByteStorable interface {
28+
~byte
29+
Storable
30+
}
31+
32+
// ByteArrayToByteSlice converts an array to []byte if the array elements are of ByteStorable type.
33+
// ByteArrayToByteSlice returns UnexpectedElementTypeError error if any array element is not of ByteStorable type.
34+
func ByteArrayToByteSlice[T ByteStorable](array *Array) ([]byte, error) {
35+
if array.Count() == 0 {
36+
return nil, nil
37+
}
38+
39+
// Get first array data slab for traversal.
40+
slab, err := firstArrayDataSlab(array.Storage, array.root)
41+
if err != nil {
42+
// Don't need to wrap error as external error because err is already categorized by firstArrayDataSlab().
43+
return nil, err
44+
}
45+
46+
data := make([]byte, 0, array.Count())
47+
48+
// Traverse array data slabs to construct []byte.
49+
for {
50+
for _, e := range slab.elements {
51+
b, ok := e.(T)
52+
if !ok {
53+
return nil, NewUnexpectedElementTypeError(reflect.TypeFor[T](), reflect.TypeOf(e))
54+
}
55+
data = append(data, byte(b))
56+
}
57+
58+
if slab.next == SlabIDUndefined {
59+
break
60+
}
61+
62+
nextSlab, err := getArraySlab(array.Storage, slab.next)
63+
if err != nil {
64+
// Don't need to wrap error as external error because err is already categorized by getArraySlab().
65+
return nil, err
66+
}
67+
68+
slab = nextSlab.(*ArrayDataSlab)
69+
}
70+
71+
return data, nil
72+
}
73+
74+
// ByteStorableValue is a type constraint that permits any type that
75+
// implements Storable and Value interfaces, and has byte underlying type.
76+
type ByteStorableValue interface {
77+
~byte
78+
Storable
79+
Value
80+
}
81+
82+
const (
83+
byteStorableCBORTagSize = 2
84+
byteStorableCBORDataSize = 2
85+
)
86+
87+
// ByteSliceToByteArray converts []byte to an array.
88+
func ByteSliceToByteArray[T ByteStorableValue](
89+
storage SlabStorage,
90+
address Address,
91+
typeInfo TypeInfo,
92+
data []byte,
93+
estimatedByteStorableSize uint32,
94+
) (*Array, error) {
95+
if len(data) == 0 {
96+
return NewArray(storage, address, typeInfo)
97+
}
98+
99+
if estimatedByteStorableSize == 0 {
100+
estimatedByteStorableSize = byteStorableCBORTagSize + byteStorableCBORDataSize
101+
}
102+
103+
estimatedEncodedDataSize := int(estimatedByteStorableSize) * len(data)
104+
105+
// If data can fit into a single data slab, create a new array with root data slab directly.
106+
if estimatedEncodedDataSize+arrayRootDataSlabPrefixSize < int(targetThreshold) {
107+
108+
elementSize := uint32(0)
109+
elements := make([]Storable, len(data))
110+
for i, b := range data {
111+
e := T(b)
112+
elements[i] = e
113+
elementSize += e.ByteSize()
114+
}
115+
116+
if elementSize+arrayRootDataSlabPrefixSize < targetThreshold {
117+
return newArrayWithElements(storage, address, typeInfo, elements, elementSize)
118+
}
119+
}
120+
121+
// Since data is too large to fit into a single data slab, call
122+
// NewArrayFromBatchData() to create an array with multiple slabs.
123+
124+
index := 0
125+
return NewArrayFromBatchData(
126+
storage,
127+
address,
128+
typeInfo,
129+
func() (Value, error) {
130+
if index == len(data) {
131+
return nil, nil
132+
}
133+
v := data[index]
134+
index++
135+
return T(v), nil
136+
})
137+
}
138+
139+
func newArrayWithElements(
140+
storage SlabStorage,
141+
address Address,
142+
typeInfo TypeInfo,
143+
elements []Storable,
144+
elementSize uint32,
145+
) (*Array, error) {
146+
// Create an empty array.
147+
array, err := NewArray(storage, address, typeInfo)
148+
if err != nil {
149+
return nil, err
150+
}
151+
152+
// Modify array root slab to include elements.
153+
root := array.root.(*ArrayDataSlab)
154+
root.elements = elements
155+
root.header.count = uint32(len(elements))
156+
root.header.size += elementSize
157+
158+
return array, nil
159+
}

‎array_conversion_test.go‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
/*
2+
* Atree - Scalable Arrays and Ordered Maps
3+
*
4+
* Copyright Flow Foundation
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the "License");
7+
* you may not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package atree_test
20+
21+
import (
22+
"math"
23+
"testing"
24+
25+
"github.com/onflow/atree"
26+
testutils "github.com/onflow/atree/test_utils"
27+
28+
"github.com/stretchr/testify/require"
29+
)
30+
31+
func TestConversionBetweenByteArrayAndByteSlice(t *testing.T) {
32+
t.Parallel()
33+
34+
typeInfo := testutils.NewSimpleTypeInfo(42)
35+
address := atree.Address{1, 2, 3, 4, 5, 6, 7, 8}
36+
37+
estimatedElementSize := testutils.Uint8Value(math.MaxUint8).ByteSize()
38+
39+
t.Run("empty byte array", func(t *testing.T) {
40+
storage := newTestPersistentStorage(t)
41+
42+
array, err := atree.NewArray(storage, address, typeInfo)
43+
require.NoError(t, err)
44+
require.Equal(t, uint64(0), array.Count())
45+
46+
// Convert atree array to []byte.
47+
data, err := atree.ByteArrayToByteSlice[testutils.Uint8Value](array)
48+
require.NoError(t, err)
49+
require.Equal(t, 0, len(data))
50+
51+
// Convert []byte to atree array.
52+
convertedArray, err := atree.ByteSliceToByteArray[testutils.Uint8Value](storage, address, typeInfo, data, estimatedElementSize)
53+
require.NoError(t, err)
54+
require.Equal(t, uint64(0), convertedArray.Count())
55+
56+
testArray(t, storage, typeInfo, address, convertedArray, nil, false, 2)
57+
})
58+
59+
t.Run("single-slab byte array", func(t *testing.T) {
60+
const arrayCount = 10
61+
62+
storage := newTestPersistentStorage(t)
63+
64+
array, err := atree.NewArray(storage, address, typeInfo)
65+
require.NoError(t, err)
66+
require.Equal(t, uint64(0), array.Count())
67+
68+
expectedByteSlice := make([]byte, arrayCount)
69+
expectedValues := make(testutils.ExpectedArrayValue, arrayCount)
70+
for i := range arrayCount {
71+
b := byte(i)
72+
expectedByteSlice[i] = b
73+
74+
v := testutils.Uint8Value(b)
75+
expectedValues[i] = testutils.Uint8Value(b)
76+
77+
err = array.Append(v)
78+
require.NoError(t, err)
79+
}
80+
require.Equal(t, uint64(arrayCount), array.Count())
81+
require.True(t, atree.GetArrayRootSlab(array).IsData())
82+
83+
// Convert atree array to []byte.
84+
data, err := atree.ByteArrayToByteSlice[testutils.Uint8Value](array)
85+
require.NoError(t, err)
86+
require.Equal(t, arrayCount, len(data))
87+
require.Equal(t, expectedByteSlice, data)
88+
89+
// Convert []byte to atree array.
90+
convertedArray, err := atree.ByteSliceToByteArray[testutils.Uint8Value](storage, address, typeInfo, data, estimatedElementSize)
91+
require.NoError(t, err)
92+
require.Equal(t, uint64(arrayCount), convertedArray.Count())
93+
94+
for i := range arrayCount {
95+
originalElement, err := array.Get(uint64(i))
96+
require.NoError(t, err)
97+
98+
convertedElement, err := convertedArray.Get(uint64(i))
99+
require.NoError(t, err)
100+
101+
require.Equal(t, originalElement, convertedElement)
102+
}
103+
104+
testArray(t, storage, typeInfo, address, convertedArray, expectedValues, false, 2)
105+
})
106+
107+
t.Run("multi-slab byte array", func(t *testing.T) {
108+
const arrayCount = 4096
109+
110+
storage := newTestPersistentStorage(t)
111+
112+
array, err := atree.NewArray(storage, address, typeInfo)
113+
require.NoError(t, err)
114+
require.Equal(t, uint64(0), array.Count())
115+
116+
expectedByteSlice := make([]byte, arrayCount)
117+
expectedValues := make(testutils.ExpectedArrayValue, arrayCount)
118+
for i := range arrayCount {
119+
b := uint8(i % 256)
120+
expectedByteSlice[i] = b
121+
122+
v := testutils.Uint8Value(b)
123+
expectedValues[i] = testutils.Uint8Value(b)
124+
125+
err = array.Append(v)
126+
require.NoError(t, err)
127+
}
128+
require.Equal(t, uint64(arrayCount), array.Count())
129+
require.False(t, atree.GetArrayRootSlab(array).IsData())
130+
131+
// Convert atree array to []byte.
132+
data, err := atree.ByteArrayToByteSlice[testutils.Uint8Value](array)
133+
require.NoError(t, err)
134+
require.Equal(t, arrayCount, len(data))
135+
require.Equal(t, expectedByteSlice, data)
136+
137+
// Convert []byte to atree array.
138+
convertedArray, err := atree.ByteSliceToByteArray[testutils.Uint8Value](storage, address, typeInfo, data, estimatedElementSize)
139+
require.NoError(t, err)
140+
require.Equal(t, uint64(arrayCount), convertedArray.Count())
141+
142+
for i := range arrayCount {
143+
originalElement, err := array.Get(uint64(i))
144+
require.NoError(t, err)
145+
146+
convertedElement, err := convertedArray.Get(uint64(i))
147+
require.NoError(t, err)
148+
149+
require.Equal(t, originalElement, convertedElement)
150+
}
151+
152+
testArray(t, storage, typeInfo, address, convertedArray, expectedValues, false, 2)
153+
})
154+
155+
t.Run("array with non byte elements", func(t *testing.T) {
156+
const arrayCount = 10
157+
158+
storage := newTestPersistentStorage(t)
159+
160+
array, err := atree.NewArray(storage, address, typeInfo)
161+
require.NoError(t, err)
162+
require.Equal(t, uint64(0), array.Count())
163+
164+
for i := range arrayCount {
165+
b := byte(i)
166+
v := testutils.Uint64Value(b)
167+
168+
err = array.Append(v)
169+
require.NoError(t, err)
170+
}
171+
require.Equal(t, uint64(arrayCount), array.Count())
172+
require.True(t, atree.GetArrayRootSlab(array).IsData())
173+
174+
_, err = atree.ByteArrayToByteSlice[testutils.Uint8Value](array)
175+
require.Error(t, err)
176+
var unexpectedElementTypeError *atree.UnexpectedElementTypeError
177+
require.ErrorAs(t, err, &unexpectedElementTypeError)
178+
require.ErrorContains(t, err, "invalid element type: expected testutils.Uint8Value, got testutils.Uint64Value")
179+
})
180+
}

0 commit comments

Comments
 (0)