Lazy RLE+

chain/actors/actor_miner.go

@@ -477,7 +477,10 @@ func (sma StorageMinerActor) SubmitFallbackPoSt(act *types.Actor, vmctx types.VM
 		return nil, aerrors.Absorb(err, 3, "could not decode sectorset")
 	}
 
-	faults := self.CurrentFaultSet.All()
+	faults, nerr := self.CurrentFaultSet.All()
+	if nerr != nil {
+		return nil, aerrors.Absorb(err, 5, "RLE+ invalid")
+	}
 	_ = faults
 
 	proverID := vmctx.Message().To // TODO: normalize to ID address

chain/types/bitfield.go

@@ -3,19 +3,27 @@ package types
 import (
 	"fmt"
 	"io"
-	"sort"
 
-	"github.com/filecoin-project/lotus/extern/rleplus"
+	rlepluslazy "github.com/filecoin-project/lotus/lib/rlepluslazy"
 	cbg "github.com/whyrusleeping/cbor-gen"
 	"golang.org/x/xerrors"
 )
 
 type BitField struct {
+	rle rlepluslazy.RLE
+
 	bits map[uint64]struct{}
 }
 
 func NewBitField() BitField {
-	return BitField{bits: make(map[uint64]struct{})}
+	rle, err := rlepluslazy.FromBuf([]byte{})
+	if err != nil {
+		panic(err)
+	}
+	return BitField{
+		rle:  rle,
+		bits: make(map[uint64]struct{}),
+	}
 }
 
 func BitFieldFromSet(setBits []uint64) BitField {
@@ -26,31 +34,59 @@ func BitFieldFromSet(setBits []uint64) BitField {
 	return res
 }
 
+func (bf BitField) sum() (rlepluslazy.RunIterator, error) {
+	if len(bf.bits) == 0 {
+		return bf.rle.RunIterator()
+	}
+
+	a, err := bf.rle.RunIterator()
+	if err != nil {
+		return nil, err
+	}
+	slc := make([]uint64, 0, len(bf.bits))
+	for b := range bf.bits {
+		slc = append(slc, b)
+	}
+
+	b, err := rlepluslazy.RunsFromSlice(slc)
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := rlepluslazy.Sum(a, b)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
 // Set ...s bit in the BitField
 func (bf BitField) Set(bit uint64) {
 	bf.bits[bit] = struct{}{}
 }
 
-// Clear ...s bit in the BitField
-func (bf BitField) Clear(bit uint64) {
-	delete(bf.bits, bit)
+func (bf BitField) Count() (uint64, error) {
+	s, err := bf.sum()
+	if err != nil {
+		return 0, err
+	}
+	return rlepluslazy.Count(s)
 }
 
-// Has checkes if bit is set in the BitField
-func (bf BitField) Has(bit uint64) bool {
-	_, ok := bf.bits[bit]
-	return ok
-}
+// All returns all set bits
+func (bf BitField) All() ([]uint64, error) {
 
-// All returns all set bits, in random order
-func (bf BitField) All() []uint64 {
-	res := make([]uint64, 0, len(bf.bits))
-	for i := range bf.bits {
-		res = append(res, i)
+	runs, err := bf.sum()
+	if err != nil {
+		return nil, err
 	}
 
-	sort.Slice(res, func(i, j int) bool { return res[i] < res[j] })
-	return res
+	res, err := rlepluslazy.SliceFromRuns(runs)
+	if err != nil {
+		return nil, err
+	}
+
+	return res, err
 }
 
 func (bf BitField) MarshalCBOR(w io.Writer) error {
@@ -59,7 +95,12 @@ func (bf BitField) MarshalCBOR(w io.Writer) error {
 		ints = append(ints, i)
 	}
 
-	rle, _, err := rleplus.Encode(ints) // Encode sorts internally
+	s, err := bf.sum()
+	if err != nil {
+		return err
+	}
+
+	rle, err := rlepluslazy.EncodeRuns(s, []byte{})
 	if err != nil {
 		return err
 	}
@@ -88,19 +129,17 @@ func (bf *BitField) UnmarshalCBOR(r io.Reader) error {
 		return fmt.Errorf("expected byte array")
 	}
 
-	rle := make([]byte, extra)
-	if _, err := io.ReadFull(br, rle); err != nil {
+	buf := make([]byte, extra)
+	if _, err := io.ReadFull(br, buf); err != nil {
 		return err
 	}
 
-	ints, err := rleplus.Decode(rle)
+	rle, err := rlepluslazy.FromBuf(buf)
 	if err != nil {
 		return xerrors.Errorf("could not decode rle+: %w", err)
 	}
+	bf.rle = rle
 	bf.bits = make(map[uint64]struct{})
-	for _, i := range ints {
-		bf.bits[i] = struct{}{}
-	}
 
 	return nil
 }

lib/rlepluslazy/bits.go

@@ -0,0 +1,160 @@
+package rlepluslazy
+
+import (
+	"sort"
+)
+
+type it2b struct {
+	source RunIterator
+	curIdx uint64
+
+	run Run
+}
+
+func (it *it2b) HasNext() bool {
+	return it.run.Valid()
+}
+
+func (it *it2b) Next() (uint64, error) {
+	it.run.Len--
+	res := it.curIdx
+	it.curIdx++
+	return res, it.prep()
+}
+
+func (it *it2b) prep() error {
+	for !it.run.Valid() && it.source.HasNext() {
+		var err error
+		it.run, err = it.source.NextRun()
+		if err != nil {
+			return err
+		}
+
+		if !it.run.Val {
+			it.curIdx += it.run.Len
+			it.run.Len = 0
+		}
+	}
+	return nil
+}
+
+func BitsFromRuns(source RunIterator) (BitIterator, error) {
+	it := &it2b{source: source}
+	if err := it.prep(); err != nil {
+		return nil, err
+	}
+	return it, nil
+}
+
+type sliceIt struct {
+	s []uint64
+}
+
+func (it sliceIt) HasNext() bool {
+	return len(it.s) != 0
+}
+
+func (it *sliceIt) Next() (uint64, error) {
+	res := it.s[0]
+	it.s = it.s[1:]
+	return res, nil
+}
+
+func BitsFromSlice(slice []uint64) BitIterator {
+	sort.Slice(slice, func(i, j int) bool { return slice[i] < slice[j] })
+	return &sliceIt{slice}
+}
+
+type it2r struct {
+	source BitIterator
+
+	runIdx uint64
+	run    [2]Run
+}
+
+func (it *it2r) HasNext() bool {
+	return it.run[0].Valid()
+}
+
+func (it *it2r) NextRun() (Run, error) {
+	res := it.run[0]
+	it.runIdx = it.runIdx + res.Len
+	it.run[0], it.run[1] = it.run[1], Run{}
+	return res, it.prep()
+}
+
+func (it *it2r) prep() error {
+	if !it.HasNext() {
+		return nil
+	}
+	if it.run[0].Val == false {
+		it.run[1].Val = true
+		it.run[1].Len = 1
+		return nil
+	}
+
+	for it.source.HasNext() && !it.run[1].Valid() {
+		nB, err := it.source.Next()
+		if err != nil {
+			return err
+		}
+
+		//fmt.Printf("runIdx: %d, run[0].Len: %d, nB: %d\n", it.runIdx, it.run[0].Len, nB)
+		if it.runIdx+it.run[0].Len == nB {
+			it.run[0].Len++
+		} else {
+			it.run[1].Len = nB - it.runIdx - it.run[0].Len
+			it.run[1].Val = false
+		}
+	}
+	return nil
+}
+
+func (it *it2r) init() error {
+	if it.source.HasNext() {
+		nB, err := it.source.Next()
+		if err != nil {
+			return err
+		}
+		it.run[0].Len = nB
+		it.run[0].Val = false
+		it.run[1].Len = 1
+		it.run[1].Val = true
+	}
+
+	if !it.run[0].Valid() {
+		it.run[0], it.run[1] = it.run[1], Run{}
+		return it.prep()
+	}
+	return nil
+}
+
+func SliceFromRuns(source RunIterator) ([]uint64, error) {
+	rit, err := BitsFromRuns(source)
+	if err != nil {
+		return nil, err
+	}
+
+	res := make([]uint64, 0)
+	for rit.HasNext() {
+		bit, err := rit.Next()
+		if err != nil {
+			return nil, err
+		}
+		res = append(res, bit)
+	}
+	return res, nil
+}
+
+func RunsFromBits(source BitIterator) (RunIterator, error) {
+	it := &it2r{source: source}
+
+	if err := it.init(); err != nil {
+		return nil, err
+	}
+	return it, nil
+}
+
+func RunsFromSlice(slice []uint64) (RunIterator, error) {
+	return RunsFromBits(BitsFromSlice(slice))
+}

lib/rlepluslazy/bits_test.go

@@ -0,0 +1,27 @@
+package rlepluslazy
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRunsFromBits(t *testing.T) {
+	expected := []Run{Run{Val: false, Len: 0x1},
+		{Val: true, Len: 0x3},
+		{Val: false, Len: 0x2},
+		{Val: true, Len: 0x3},
+	}
+	rit, err := RunsFromBits(BitsFromSlice([]uint64{1, 2, 3, 6, 7, 8}))
+	assert.NoError(t, err)
+	i := 10
+	output := make([]Run, 0, 4)
+	for rit.HasNext() && i > 0 {
+		run, err := rit.NextRun()
+		assert.NoError(t, err)
+		i--
+		output = append(output, run)
+	}
+	assert.NotEqual(t, 0, i, "too many iterations")
+	assert.Equal(t, expected, output)
+}