wehub-resource-sync
354 lines
8.7 KiB
Go
354 lines
8.7 KiB
Go
// Copyright 2021 Dolthub, Inc.
|
|
//
|
|
// 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 tree
|
|
|
|
import (
|
|
"context"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"sync/atomic"
|
|
|
|
"github.com/dolthub/dolt/go/gen/fb/serial"
|
|
"github.com/dolthub/dolt/go/libraries/doltcore/schema"
|
|
"github.com/dolthub/dolt/go/store/chunks"
|
|
"github.com/dolthub/dolt/go/store/hash"
|
|
"github.com/dolthub/dolt/go/store/prolly/message"
|
|
"github.com/dolthub/dolt/go/store/types"
|
|
"github.com/dolthub/dolt/go/store/val"
|
|
)
|
|
|
|
type Item []byte
|
|
|
|
type subtreeCounts []uint64
|
|
|
|
// Node is a generic implementation of a prolly tree node.
|
|
// Elements in a Node are generic Items. Interpreting Item
|
|
// contents is deferred to higher layers (see prolly.Map).
|
|
type Node struct {
|
|
// nodeHash contains the calculated hash.Hash
|
|
// for this node
|
|
nodeHash atomic.Value
|
|
// subtrees contains the key cardinality
|
|
// of each child tree of a non-leaf Node.
|
|
// this field is lazily decoded from msg
|
|
// because it requires a malloc.
|
|
subtrees atomic.Pointer[subtreeCounts]
|
|
// keys and values cache offset metadata
|
|
// to accelerate Item lookups into msg.
|
|
keys *message.ItemAccess
|
|
values *message.ItemAccess
|
|
// msg is the underlying buffer for the Node
|
|
// encoded as a Flatbuffers message.
|
|
msg serial.Message
|
|
// count is the Item pair count.
|
|
count uint16
|
|
// level is 0-indexed tree height.
|
|
level uint16
|
|
}
|
|
|
|
type AddressCb func(ctx context.Context, addr hash.Hash) error
|
|
|
|
func WalkAddresses(ctx context.Context, nd *Node, ns NodeStore, cb AddressCb) error {
|
|
return walkAddresses(ctx, nd, func(ctx context.Context, addr hash.Hash) error {
|
|
if err := cb(ctx, addr); err != nil {
|
|
return err
|
|
}
|
|
|
|
if nd.IsLeaf() {
|
|
return nil
|
|
}
|
|
|
|
child, err := ns.Read(ctx, addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return WalkAddresses(ctx, child, ns, cb)
|
|
})
|
|
}
|
|
|
|
type NodeCb func(ctx context.Context, nd *Node) error
|
|
|
|
// WalkNodes runs a callback function on every node found in the DFS of |nd|
|
|
// that is of the same message type as |nd|.
|
|
func WalkNodes(ctx context.Context, nd *Node, ns NodeStore, cb NodeCb) error {
|
|
if err := cb(ctx, nd); err != nil {
|
|
return err
|
|
}
|
|
if nd == nil {
|
|
return nil
|
|
}
|
|
if nd.IsLeaf() {
|
|
return nil
|
|
}
|
|
|
|
return walkAddresses(ctx, nd, func(ctx context.Context, addr hash.Hash) error {
|
|
child, err := ns.Read(ctx, addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return WalkNodes(ctx, child, ns, cb)
|
|
})
|
|
}
|
|
|
|
// walkOpaqueNodes runs a callback function on every node found in the DFS of |nd|
|
|
// including nested trees.
|
|
func walkOpaqueNodes(ctx context.Context, nd *Node, ns NodeStore, cb NodeCb) error {
|
|
if err := cb(ctx, nd); err != nil {
|
|
return err
|
|
}
|
|
|
|
return walkAddresses(ctx, nd, func(ctx context.Context, addr hash.Hash) error {
|
|
child, err := ns.Read(ctx, addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return WalkNodes(ctx, child, ns, cb)
|
|
})
|
|
}
|
|
|
|
func NodeFromBytes(msg []byte) (node *Node, fileId string, err error) {
|
|
fileId, keys, values, level, count, err := message.UnpackFields(msg)
|
|
return &Node{
|
|
keys: keys,
|
|
values: values,
|
|
count: count,
|
|
level: level,
|
|
msg: msg,
|
|
}, fileId, err
|
|
}
|
|
|
|
func NodeFromBytesWithHash(msg []byte, hash hash.Hash) (node *Node, fileId string, err error) {
|
|
fileId, keys, values, level, count, err := message.UnpackFields(msg)
|
|
node = &Node{
|
|
keys: keys,
|
|
values: values,
|
|
count: count,
|
|
level: level,
|
|
msg: msg,
|
|
}
|
|
node.nodeHash.Store(hash)
|
|
return node, fileId, err
|
|
}
|
|
|
|
func NodeFromChunk(chunk *chunks.Chunk) (node *Node, fieldId string, err error) {
|
|
fieldId, keys, values, level, count, err := message.UnpackFields(chunk.Data())
|
|
node = &Node{
|
|
keys: keys,
|
|
values: values,
|
|
count: count,
|
|
level: level,
|
|
msg: chunk.Data(),
|
|
}
|
|
node.nodeHash.Store(chunk.Hash())
|
|
return node, fieldId, err
|
|
}
|
|
|
|
func (nd *Node) HashOf() hash.Hash {
|
|
h := nd.nodeHash.Load()
|
|
if h == nil {
|
|
h = hash.Of(nd.bytes())
|
|
nd.nodeHash.Store(h)
|
|
}
|
|
return h.(hash.Hash)
|
|
}
|
|
|
|
func (nd *Node) Count() int {
|
|
if nd == nil {
|
|
return 0
|
|
}
|
|
return int(nd.count)
|
|
}
|
|
|
|
func (nd *Node) TreeCount() (int, error) {
|
|
return message.GetTreeCount(nd.msg)
|
|
}
|
|
|
|
func (nd *Node) Size() int {
|
|
return len(nd.bytes())
|
|
}
|
|
|
|
// Level returns the tree Level for this node
|
|
func (nd *Node) Level() int {
|
|
if nd == nil {
|
|
return 0
|
|
}
|
|
return int(nd.level)
|
|
}
|
|
|
|
// IsLeaf returns whether this node is a leaf
|
|
func (nd *Node) IsLeaf() bool {
|
|
return nd.level == 0
|
|
}
|
|
|
|
// GetKey returns the |ith| key of this node
|
|
func (nd *Node) GetKey(i int) Item {
|
|
return nd.keys.GetItem(i, nd.msg)
|
|
}
|
|
|
|
// GetValue returns the |ith| value of this node.
|
|
func (nd *Node) GetValue(i int) Item {
|
|
return nd.values.GetItem(i, nd.msg)
|
|
}
|
|
|
|
func (nd *Node) LoadSubtrees() (*Node, error) {
|
|
if nd.subtrees.Load() != nil {
|
|
return nd, nil
|
|
}
|
|
// deserializing subtree counts requires a malloc,
|
|
// we don't load them unless explicitly requested
|
|
sc, err := message.GetSubtrees(nd.msg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
stc := subtreeCounts(sc)
|
|
nd.subtrees.CompareAndSwap(nil, &stc)
|
|
return nd, nil
|
|
}
|
|
|
|
func (nd *Node) GetSubtreeCount(i int) uint64 {
|
|
if nd.IsLeaf() {
|
|
return 1
|
|
}
|
|
// this will panic unless subtrees were loaded.
|
|
return (*nd.subtrees.Load())[i]
|
|
}
|
|
|
|
// getAddress returns the |ith| address of this node.
|
|
// This method assumes values are 20-byte address hashes.
|
|
func (nd *Node) getAddress(i int) hash.Hash {
|
|
return hash.New(nd.GetValue(i))
|
|
}
|
|
|
|
func (nd *Node) empty() bool {
|
|
return nd == nil || nd.bytes() == nil || nd.count == 0
|
|
}
|
|
|
|
func (nd *Node) bytes() []byte {
|
|
if nd == nil {
|
|
return nil
|
|
}
|
|
return nd.msg
|
|
}
|
|
|
|
func walkAddresses(ctx context.Context, nd *Node, cb AddressCb) (err error) {
|
|
return message.WalkAddresses(ctx, nd.msg, cb)
|
|
}
|
|
|
|
func getLastKey(nd *Node) Item {
|
|
return nd.GetKey(nd.Count() - 1)
|
|
}
|
|
|
|
// OutputProllyNode writes the node given to the writer given in a human-readable format, with values converted
|
|
// to the type specified by the provided schema. All nodes have keys displayed in this manner. Interior nodes have
|
|
// their child hash references spelled out, leaf nodes have value tuples delineated like the keys
|
|
func OutputProllyNode(ctx context.Context, w io.Writer, node *Node, ns NodeStore, schema schema.Schema) error {
|
|
kd := schema.GetKeyDescriptor(ns)
|
|
vd := schema.GetValueDescriptor(ns)
|
|
for i := 0; i < node.Count(); i++ {
|
|
k := node.GetKey(i)
|
|
kt := val.Tuple(k)
|
|
|
|
w.Write([]byte("\n { key: "))
|
|
for j := 0; j < kt.Count(); j++ {
|
|
if j > 0 {
|
|
w.Write([]byte(", "))
|
|
}
|
|
|
|
isAddr := j < len(kd.Types) && val.IsAddrEncoding(kd.Types[j].Enc)
|
|
if isAddr {
|
|
w.Write([]byte("#"))
|
|
}
|
|
w.Write([]byte(hex.EncodeToString(kd.GetField(j, kt))))
|
|
if isAddr {
|
|
w.Write([]byte(" ("))
|
|
key, err := GetField(ctx, kd, j, kt, ns)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
w.Write([]byte(fmt.Sprint(key)))
|
|
w.Write([]byte(")"))
|
|
}
|
|
|
|
}
|
|
|
|
if node.IsLeaf() {
|
|
v := node.GetValue(i)
|
|
vt := val.Tuple(v)
|
|
|
|
w.Write([]byte(" value: "))
|
|
for j := 0; j < vt.Count(); j++ {
|
|
if j > 0 {
|
|
w.Write([]byte(", "))
|
|
}
|
|
isAddr := val.IsAddrEncoding(vd.Types[j].Enc)
|
|
if isAddr {
|
|
w.Write([]byte("#"))
|
|
}
|
|
w.Write([]byte(hex.EncodeToString(vd.GetField(j, vt))))
|
|
if isAddr {
|
|
w.Write([]byte(" ("))
|
|
value, err := GetField(ctx, vd, j, vt, ns)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
w.Write([]byte(fmt.Sprint(value)))
|
|
w.Write([]byte(")"))
|
|
}
|
|
}
|
|
|
|
w.Write([]byte(" }"))
|
|
} else {
|
|
ref := node.getAddress(i)
|
|
|
|
w.Write([]byte(" ref: #"))
|
|
w.Write([]byte(ref.String()))
|
|
w.Write([]byte(" }"))
|
|
}
|
|
}
|
|
|
|
w.Write([]byte("\n"))
|
|
return nil
|
|
}
|
|
|
|
// OutputProllyNodeBytes writes the node given to the writer given in a semi-human-readable format, where values are still
|
|
// displayed in hex-encoded byte strings, but are delineated into their fields. All nodes have keys displayed in this
|
|
// manner. Interior nodes have their child hash references spelled out, leaf nodes have value tuples delineated like
|
|
// the keys
|
|
func OutputProllyNodeBytes(w io.Writer, node *Node) error {
|
|
return types.OutputProllyNodeBytes(w, node.msg)
|
|
}
|
|
|
|
func OutputAddressMapNode(w io.Writer, node *Node) error {
|
|
for i := 0; i < int(node.Count()); i++ {
|
|
k := node.GetKey(i)
|
|
w.Write([]byte("\n { key: "))
|
|
w.Write(k)
|
|
|
|
ref := node.getAddress(i)
|
|
|
|
w.Write([]byte(" ref: #"))
|
|
w.Write([]byte(ref.String()))
|
|
w.Write([]byte(" }"))
|
|
}
|
|
w.Write([]byte("\n"))
|
|
return nil
|
|
}
|
|
|
|
func ValueFromNode(root *Node) types.Value {
|
|
return types.SerialMessage(root.bytes())
|
|
}
|