dolt/walk/walk.go

package walk

import (
	"fmt"
	"sync"

	"github.com/attic-labs/noms/chunks"
	"github.com/attic-labs/noms/d"
	"github.com/attic-labs/noms/ref"
	"github.com/attic-labs/noms/types"
)

// SomeCallback takes a ref.Ref and returns a bool indicating whether
// the current walk should skip the tree descending from value.
type SomeCallback func(r ref.Ref) bool

// AllCallback takes a ref and processes it.
type AllCallback func(r ref.Ref)

// Some recursively walks over all ref.Refs reachable from r and calls cb on them. If cb ever returns true, the walk will stop recursing on the current ref. If |concurrency| > 1, it is the callers responsibility to make ensure that |cb| is threadsafe.
func SomeP(r ref.Ref, cs chunks.ChunkSource, cb SomeCallback, concurrency int) {
	doTreeWalkP(r, cs, cb, concurrency)
}

// All recursively walks over all ref.Refs reachable from r and calls cb on them. If |concurrency| > 1, it is the callers responsibility to make ensure that |cb| is threadsafe.
func AllP(r ref.Ref, cs chunks.ChunkSource, cb AllCallback, concurrency int) {
	doTreeWalkP(r, cs, func(r ref.Ref) (skip bool) {
		cb(r)
		return
	}, concurrency)
}

func doTreeWalkP(r ref.Ref, cs chunks.ChunkSource, cb SomeCallback, concurrency int) {
	rq := newRefQueue()
	f := newFailure()

	visited := map[ref.Ref]bool{}
	mu := sync.Mutex{}
	wg := sync.WaitGroup{}

	processRef := func(r ref.Ref) {
		defer wg.Done()

		mu.Lock()
		skip := cb(r) || visited[r]
		visited[r] = true
		mu.Unlock()

		if skip || f.didFail() {
			return
		}

		v := types.ReadValue(r, cs)
		if v == nil {
			f.fail(fmt.Errorf("Attempt to copy absent ref:%s", r.String()))
			return
		}

		for _, cf := range v.Chunks() {
			wg.Add(1)
			rq.tail() <- cf.Ref()
		}
	}

	iter := func() {
		for r := range rq.head() {
			processRef(r)
		}
	}

	for i := 0; i < concurrency; i++ {
		go iter()
	}

	wg.Add(1)
	rq.tail() <- r
	wg.Wait()

	rq.close()

	f.checkNotFailed()
}

// refQueue emulates a buffered channel of refs of unlimited size.
type refQueue struct {
	head  func() <-chan ref.Ref
	tail  func() chan<- ref.Ref
	close func()
}

func newRefQueue() refQueue {
	head := make(chan ref.Ref, 64)
	tail := make(chan ref.Ref, 64)
	done := make(chan struct{})
	buff := []ref.Ref{}

	push := func(r ref.Ref) {
		buff = append(buff, r)
	}

	pop := func() ref.Ref {
		d.Chk.True(len(buff) > 0)
		r := buff[0]
		buff = buff[1:]
		return r
	}

	go func() {
	loop:
		for {
			if len(buff) == 0 {
				select {
				case r := <-tail:
					push(r)
				case <-done:
					break loop
				}
			} else {
				first := buff[0]
				select {
				case r := <-tail:
					push(r)
				case head <- first:
					r := pop()
					d.Chk.Equal(r, first)
				case <-done:
					break loop
				}
			}
		}
	}()

	return refQueue{
		func() <-chan ref.Ref {
			return head
		},
		func() chan<- ref.Ref {
			return tail
		},
		func() {
			close(head)
			done <- struct{}{}
		},
	}
}

type failure struct {
	err error
	mu  *sync.Mutex
}

func newFailure() *failure {
	return &failure{
		mu: &sync.Mutex{},
	}
}

func (f *failure) fail(err error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.err == nil { // only capture first error
		f.err = err
	}
}

func (f *failure) didFail() bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.err != nil
}

func (f *failure) checkNotFailed() {
	f.mu.Lock()
	defer f.mu.Unlock()
	d.Chk.NoError(f.err)
}