github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/store/prolly/message/message.go (about)

     1  // Copyright 2022 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package message
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"math"
    21  
    22  	fb "github.com/dolthub/flatbuffers/v23/go"
    23  
    24  	"github.com/dolthub/dolt/go/gen/fb/serial"
    25  	"github.com/dolthub/dolt/go/store/hash"
    26  )
    27  
    28  const (
    29  	maxChunkSz = math.MaxUint16
    30  	addrSize   = hash.ByteLen
    31  	uint16Size = 2
    32  )
    33  
    34  type Serializer interface {
    35  	Serialize(keys, values [][]byte, subtrees []uint64, level int) serial.Message
    36  }
    37  
    38  func UnpackFields(msg serial.Message) (keys, values ItemAccess, level, count uint16, err error) {
    39  	switch serial.GetFileID(msg) {
    40  	case serial.ProllyTreeNodeFileID:
    41  		return getProllyMapKeysAndValues(msg)
    42  	case serial.AddressMapFileID:
    43  		keys, err = getAddressMapKeys(msg)
    44  		if err != nil {
    45  			return
    46  		}
    47  		values, err = getAddressMapValues(msg)
    48  		if err != nil {
    49  			return
    50  		}
    51  		level, err = getAddressMapTreeLevel(msg)
    52  		if err != nil {
    53  			return
    54  		}
    55  		count, err = getAddressMapCount(msg)
    56  		return
    57  	case serial.MergeArtifactsFileID:
    58  		return getArtifactMapKeysAndValues(msg)
    59  	case serial.CommitClosureFileID:
    60  		keys, err = getCommitClosureKeys(msg)
    61  		if err != nil {
    62  			return
    63  		}
    64  		values, err = getCommitClosureValues(msg)
    65  		if err != nil {
    66  			return
    67  		}
    68  		level, err = getCommitClosureTreeLevel(msg)
    69  		if err != nil {
    70  			return
    71  		}
    72  		count, err = getCommitClosureCount(msg)
    73  		return
    74  	case serial.BlobFileID:
    75  		keys, err = getBlobKeys(msg)
    76  		if err != nil {
    77  			return
    78  		}
    79  		values, err = getBlobValues(msg)
    80  		if err != nil {
    81  			return
    82  		}
    83  		level, err = getBlobTreeLevel(msg)
    84  		if err != nil {
    85  			return
    86  		}
    87  		count, err = getBlobCount(msg)
    88  		return
    89  	default:
    90  		panic(fmt.Sprintf("unknown message id %s", serial.GetFileID(msg)))
    91  	}
    92  }
    93  
    94  func WalkAddresses(ctx context.Context, msg serial.Message, cb func(ctx context.Context, addr hash.Hash) error) error {
    95  	id := serial.GetFileID(msg)
    96  	switch id {
    97  	case serial.ProllyTreeNodeFileID:
    98  		return walkProllyMapAddresses(ctx, msg, cb)
    99  	case serial.AddressMapFileID:
   100  		return walkAddressMapAddresses(ctx, msg, cb)
   101  	case serial.MergeArtifactsFileID:
   102  		return walkMergeArtifactAddresses(ctx, msg, cb)
   103  	case serial.CommitClosureFileID:
   104  		return walkCommitClosureAddresses(ctx, msg, cb)
   105  	case serial.BlobFileID:
   106  		return walkBlobAddresses(ctx, msg, cb)
   107  	default:
   108  		panic(fmt.Sprintf("unknown message id %s", id))
   109  	}
   110  }
   111  
   112  func GetTreeCount(msg serial.Message) (int, error) {
   113  	id := serial.GetFileID(msg)
   114  	switch id {
   115  	case serial.ProllyTreeNodeFileID:
   116  		return getProllyMapTreeCount(msg)
   117  	case serial.AddressMapFileID:
   118  		return getAddressMapTreeCount(msg)
   119  	case serial.MergeArtifactsFileID:
   120  		return getMergeArtifactTreeCount(msg)
   121  	case serial.CommitClosureFileID:
   122  		return getCommitClosureTreeCount(msg)
   123  	case serial.BlobFileID:
   124  		return getBlobTreeCount(msg)
   125  	default:
   126  		panic(fmt.Sprintf("unknown message id %s", id))
   127  	}
   128  }
   129  
   130  func GetSubtrees(msg serial.Message) ([]uint64, error) {
   131  	id := serial.GetFileID(msg)
   132  	switch id {
   133  	case serial.ProllyTreeNodeFileID:
   134  		return getProllyMapSubtrees(msg)
   135  	case serial.AddressMapFileID:
   136  		return getAddressMapSubtrees(msg)
   137  	case serial.MergeArtifactsFileID:
   138  		return getMergeArtifactSubtrees(msg)
   139  	case serial.CommitClosureFileID:
   140  		return getCommitClosureSubtrees(msg)
   141  	case serial.BlobFileID:
   142  		return getBlobSubtrees(msg)
   143  	default:
   144  		panic(fmt.Sprintf("unknown message id %s", id))
   145  	}
   146  }
   147  
   148  func lookupVectorOffset(vo fb.VOffsetT, tab fb.Table) uint16 {
   149  	off := fb.UOffsetT(tab.Offset(vo)) + tab.Pos
   150  	off += fb.GetUOffsetT(tab.Bytes[off:])
   151  	// data starts after metadata containing the vector length
   152  	return uint16(off + fb.UOffsetT(fb.SizeUOffsetT))
   153  }
   154  
   155  func assertTrue(b bool, msg string) {
   156  	if !b {
   157  		panic("assertion failed: " + msg)
   158  	}
   159  }
   160  
   161  func assertFalse(b bool, msg string) {
   162  	if b {
   163  		panic("assertion failed: " + msg)
   164  	}
   165  }