Merge pull request #508 from onflow/fxamacker/reorganize-array-code

Split array.go into smaller files and group related functions

array_data_slab_encode.go

@@ -0,0 +1,249 @@
+/*
+ * Atree - Scalable Arrays and Ordered Maps
+ *
+ * Copyright Flow Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package atree
+
+import (
+	"encoding/binary"
+	"fmt"
+)
+
+// encodeAsInlined encodes inlined array data slab. Encoding is
+// version 1 with CBOR tag having tag number CBORTagInlinedArray,
+// and tag contant as 3-element array:
+//
+//	+------------------+----------------+----------+
+//	| extra data index | value ID index | elements |
+//	+------------------+----------------+----------+
+func (a *ArrayDataSlab) encodeAsInlined(enc *Encoder) error {
+	if a.extraData == nil {
+		return NewEncodingError(
+			fmt.Errorf("failed to encode non-root array data slab as inlined"))
+	}
+
+	if !a.inlined {
+		return NewEncodingError(
+			fmt.Errorf("failed to encode standalone array data slab as inlined"))
+	}
+
+	extraDataIndex, err := enc.inlinedExtraData().addArrayExtraData(a.extraData)
+	if err != nil {
+		// err is already categorized by InlinedExtraData.addArrayExtraData().
+		return err
+	}
+
+	if extraDataIndex > maxInlinedExtraDataIndex {
+		return NewEncodingError(
+			fmt.Errorf("failed to encode inlined array data slab: extra data index %d exceeds limit %d", extraDataIndex, maxInlinedExtraDataIndex))
+	}
+
+	// Encode tag number and array head of 3 elements
+	err = enc.CBOR.EncodeRawBytes([]byte{
+		// tag number
+		0xd8, CBORTagInlinedArray,
+		// array head of 3 elements
+		0x83,
+	})
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	// element 0: extra data index
+	// NOTE: encoded extra data index is fixed sized CBOR uint
+	err = enc.CBOR.EncodeRawBytes([]byte{
+		0x18,
+		byte(extraDataIndex),
+	})
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	// element 1: slab index
+	err = enc.CBOR.EncodeBytes(a.header.slabID.index[:])
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	// element 2: array elements
+	err = a.encodeElements(enc)
+	if err != nil {
+		// err is already categorized by ArrayDataSlab.encodeElements().
+		return err
+	}
+
+	err = enc.CBOR.Flush()
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	return nil
+}
+
+// Encode encodes this array data slab to the given encoder.
+//
+// DataSlab Header:
+//
+//	+-------------------------------+----------------------+---------------------------------+-----------------------------+
+//	| slab version + flag (2 bytes) | extra data (if root) | inlined extra data (if present) | next slab ID (if non-empty) |
+//	+-------------------------------+----------------------+---------------------------------+-----------------------------+
+//
+// Content:
+//
+//	CBOR encoded array of elements
+//
+// See ArrayExtraData.Encode() for extra data section format.
+// See InlinedExtraData.Encode() for inlined extra data section format.
+func (a *ArrayDataSlab) Encode(enc *Encoder) error {
+
+	if a.inlined {
+		return a.encodeAsInlined(enc)
+	}
+
+	// Encoding is done in two steps:
+	//
+	// 1. Encode array elements using a new buffer while collecting inlined extra data from inlined elements.
+	// 2. Encode slab with deduplicated inlined extra data and copy encoded elements from previous buffer.
+
+	// Get a buffer from a pool to encode elements.
+	elementBuf := getBuffer()
+	defer putBuffer(elementBuf)
+
+	elementEnc := NewEncoder(elementBuf, enc.encMode)
+
+	err := a.encodeElements(elementEnc)
+	if err != nil {
+		// err is already categorized by Array.encodeElements().
+		return err
+	}
+
+	err = elementEnc.CBOR.Flush()
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	const version = 1
+
+	h, err := newArraySlabHead(version, slabArrayData)
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	if a.HasPointer() {
+		h.setHasPointers()
+	}
+
+	if a.next != SlabIDUndefined {
+		h.setHasNextSlabID()
+	}
+
+	if a.extraData != nil {
+		h.setRoot()
+	}
+
+	if elementEnc.hasInlinedExtraData() {
+		h.setHasInlinedSlabs()
+	}
+
+	// Encode head (version + flag)
+	_, err = enc.Write(h[:])
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	// Encode extra data
+	if a.extraData != nil {
+		// Use defaultEncodeTypeInfo to encode root level TypeInfo as is.
+		err = a.extraData.Encode(enc, defaultEncodeTypeInfo)
+		if err != nil {
+			// err is already categorized by ArrayExtraData.Encode().
+			return err
+		}
+	}
+
+	// Encode inlined extra data
+	if elementEnc.hasInlinedExtraData() {
+		err = elementEnc.inlinedExtraData().Encode(enc)
+		if err != nil {
+			// err is already categorized by inlinedExtraData.Encode().
+			return err
+		}
+	}
+
+	// Encode next slab ID
+	if a.next != SlabIDUndefined {
+		n, err := a.next.ToRawBytes(enc.Scratch[:])
+		if err != nil {
+			// Don't need to wrap because err is already categorized by SlabID.ToRawBytes().
+			return err
+		}
+
+		_, err = enc.Write(enc.Scratch[:n])
+		if err != nil {
+			return NewEncodingError(err)
+		}
+	}
+
+	// Encode elements by copying raw bytes from previous buffer
+	err = enc.CBOR.EncodeRawBytes(elementBuf.Bytes())
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	err = enc.CBOR.Flush()
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	return nil
+}
+
+func (a *ArrayDataSlab) encodeElements(enc *Encoder) error {
+	// Encode CBOR array size manually for fix-sized encoding
+
+	enc.Scratch[0] = 0x80 | 25
+
+	countOffset := 1
+	const countSize = 2
+	binary.BigEndian.PutUint16(
+		enc.Scratch[countOffset:],
+		uint16(len(a.elements)),
+	)
+
+	// Write scratch content to encoder
+	totalSize := countOffset + countSize
+	err := enc.CBOR.EncodeRawBytes(enc.Scratch[:totalSize])
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	// Encode data slab content (array of elements)
+	for _, e := range a.elements {
+		err = e.Encode(enc)
+		if err != nil {
+			// Wrap err as external error (if needed) because err is returned by Storable interface.
+			return wrapErrorfAsExternalErrorIfNeeded(err, "failed to encode array element")
+		}
+	}
+
+	err = enc.CBOR.Flush()
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	return nil
+}

array_extradata.go

@@ -0,0 +1,110 @@
+/*
+ * Atree - Scalable Arrays and Ordered Maps
+ *
+ * Copyright Flow Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package atree
+
+import (
+	"fmt"
+
+	"github.com/fxamacker/cbor/v2"
+)
+
+type ArrayExtraData struct {
+	TypeInfo TypeInfo // array type
+}
+
+var _ ExtraData = &ArrayExtraData{}
+
+const arrayExtraDataLength = 1
+
+func newArrayExtraDataFromData(
+	data []byte,
+	decMode cbor.DecMode,
+	decodeTypeInfo TypeInfoDecoder,
+) (
+	*ArrayExtraData,
+	[]byte,
+	error,
+) {
+	dec := decMode.NewByteStreamDecoder(data)
+
+	extraData, err := newArrayExtraData(dec, decodeTypeInfo)
+	if err != nil {
+		return nil, data, err
+	}
+
+	return extraData, data[dec.NumBytesDecoded():], nil
+}
+
+// newArrayExtraData decodes CBOR array to extra data:
+//
+//	cborArray{type info}
+func newArrayExtraData(dec *cbor.StreamDecoder, decodeTypeInfo TypeInfoDecoder) (*ArrayExtraData, error) {
+	length, err := dec.DecodeArrayHead()
+	if err != nil {
+		return nil, NewDecodingError(err)
+	}
+
+	if length != arrayExtraDataLength {
+		return nil, NewDecodingError(
+			fmt.Errorf(
+				"array extra data has invalid length %d, want %d",
+				length,
+				arrayExtraDataLength,
+			))
+	}
+
+	typeInfo, err := decodeTypeInfo(dec)
+	if err != nil {
+		// Wrap err as external error (if needed) because err is returned by TypeInfoDecoder callback.
+		return nil, wrapErrorfAsExternalErrorIfNeeded(err, "failed to decode type info")
+	}
+
+	return &ArrayExtraData{TypeInfo: typeInfo}, nil
+}
+
+// Encode encodes extra data as CBOR array:
+//
+//	[type info]
+func (a *ArrayExtraData) Encode(enc *Encoder, encodeTypeInfo encodeTypeInfo) error {
+	err := enc.CBOR.EncodeArrayHead(arrayExtraDataLength)
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	err = encodeTypeInfo(enc, a.TypeInfo)
+	if err != nil {
+		// Wrap err as external error (if needed) because err is returned by TypeInfo interface.
+		return wrapErrorfAsExternalErrorIfNeeded(err, "failed to encode type info")
+	}
+
+	err = enc.CBOR.Flush()
+	if err != nil {
+		return NewEncodingError(err)
+	}
+
+	return nil
+}
+
+func (a *ArrayExtraData) isExtraData() bool {
+	return true
+}
+
+func (a *ArrayExtraData) Type() TypeInfo {
+	return a.TypeInfo
+}