github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/types/walk.go (about)

     1  // Copyright 2019 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  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package types
    23  
    24  import (
    25  	"context"
    26  
    27  	"github.com/dolthub/dolt/go/store/hash"
    28  )
    29  
    30  type SkipValueCallback func(v Value) bool
    31  
    32  // WalkValues loads prolly trees progressively by walking down the tree. We don't wants to invoke
    33  // the value callback on internal sub-trees (which are valid values) because they are not logical
    34  // values in the graph
    35  type valueRec struct {
    36  	v  Value
    37  	cb bool
    38  }
    39  
    40  const maxRefCount = 1 << 12 // ~16MB of data
    41  
    42  // WalkValues recursively walks over all types.Values reachable from r and calls cb on them.
    43  func WalkValues(ctx context.Context, nbf *NomsBinFormat, target Value, vr ValueReader, cb SkipValueCallback) error {
    44  	visited := hash.HashSet{}
    45  	refs := map[hash.Hash]bool{}
    46  	values := []valueRec{{target, true}}
    47  
    48  	for len(values) > 0 || len(refs) > 0 {
    49  		for len(values) > 0 {
    50  			rec := values[len(values)-1]
    51  			values = values[:len(values)-1]
    52  
    53  			v := rec.v
    54  			if rec.cb && cb(v) {
    55  				continue
    56  			}
    57  
    58  			if _, ok := v.(Blob); ok {
    59  				continue // don't traverse into blob ptrees
    60  			}
    61  
    62  			if r, ok := v.(Ref); ok {
    63  				refs[r.TargetHash()] = true
    64  				continue
    65  			}
    66  
    67  			if col, ok := v.(Collection); ok && !col.asSequence().isLeaf() {
    68  				err := col.WalkRefs(nbf, func(r Ref) error {
    69  					refs[r.TargetHash()] = false
    70  					return nil
    71  				})
    72  
    73  				if err != nil {
    74  					return err
    75  				}
    76  
    77  				continue
    78  			}
    79  
    80  			err := v.WalkValues(ctx, func(sv Value) error {
    81  				values = append(values, valueRec{sv, true})
    82  
    83  				return nil
    84  			})
    85  
    86  			if err != nil {
    87  				return err
    88  			}
    89  		}
    90  
    91  		if len(refs) == 0 {
    92  			continue
    93  		}
    94  
    95  		hs := make(hash.HashSlice, 0, len(refs))
    96  		oldRefs := refs
    97  		refs = map[hash.Hash]bool{}
    98  		for h := range oldRefs {
    99  			if _, ok := visited[h]; ok {
   100  				continue
   101  			}
   102  
   103  			if len(hs) >= maxRefCount {
   104  				refs[h] = oldRefs[h]
   105  				continue
   106  			}
   107  
   108  			hs = append(hs, h)
   109  			visited.Insert(h)
   110  		}
   111  
   112  		if len(hs) > 0 {
   113  			readValues, err := vr.ReadManyValues(ctx, hs)
   114  
   115  			if err != nil {
   116  				return err
   117  			}
   118  
   119  			for i, sv := range readValues {
   120  				values = append(values, valueRec{sv, oldRefs[hs[i]]})
   121  			}
   122  		}
   123  	}
   124  
   125  	return nil
   126  }