Merge pull request #1097 from jaten-molecula/bitmap_foreach

rbf: BitN needs int32 to hold its maximum value.

rbf/cursor.go

@@ -350,11 +350,13 @@ func toPgno(val []byte) uint32 {
 	return binary.LittleEndian.Uint32(val)
 }
 func (c *Cursor) putLeafCell(in leafCell) (err error) {
+
 	leafPage := c.leafPage // the last read leaf page
 	cells := readLeafCells(leafPage, c.leafCells[:])
 	elem := &c.stack.elems[c.stack.index]
 	cell := in
 	if elem.index >= len(cells) || c.Key() != cell.Key {
+
 		//new cell
 		if in.Type == ContainerTypeBitmap {
 			//allocated bitmap()
@@ -403,7 +405,6 @@ func (c *Cursor) putLeafCell(in leafCell) (err error) {
 		}
 		cell.Data = fromPgno(bitmapPgno)
 	}
-
 	cells[elem.index] = cell
 
 	// Split into multiple pages if page size is exceeded.

rbf/cursor_test.go

@@ -141,7 +141,7 @@ func TestCursor_FirstNext_Quick(t *testing.T) {
 			} else if got, want := c.Key(), item.key; got != want {
 				t.Fatalf("Key()=%d, want %d", got, want)
 			} else if got, want := c.Values(), item.values; !reflect.DeepEqual(got, want) {
-				t.Fatalf("len(Values())=%v, want %v", len(got), len(want))
+				t.Fatalf("len(Values())=%v/%#v, want %v/%#v", len(got), got, len(want), want)
 			}
 		}
 	})
@@ -707,7 +707,7 @@ func TestCursor_BitmapToArrayConversion(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	// start with n = 4084 bits, where ArrayMaxSize =
+	// start with n = 4084 bits, over ArrayMaxSize
 	rb := roaring.NewBitmap()
 	bits := make([]uint64, rbf.BitmapN)
 	n := 0
@@ -718,7 +718,6 @@ func TestCursor_BitmapToArrayConversion(t *testing.T) {
 			break
 		}
 	}
-	//vv("ArrayMaxSize=%v;  n = %v", rbf.ArrayMaxSize, n)
 	rb.Put(0, roaring.NewContainerBitmap(n, bits))
 
 	// setup to delete random bits down
@@ -1116,3 +1115,82 @@ func TestCursor_GenerateAll(t *testing.T) {
 		panic(err)
 	}
 }
+
+// test ForEachRange handles Bitmaps, because BitmapPtr
+// case was missing
+func TestForEachRange(t *testing.T) {
+	db := MustOpenDB(t)
+	defer MustCloseDB(t, db)
+	tx := MustBegin(t, db, true)
+	defer tx.Rollback()
+
+	if err := tx.CreateBitmap("x"); err != nil {
+		t.Fatal(err)
+	}
+
+	rb := roaring.NewBitmap()
+	bits := make([]uint64, rbf.BitmapN)
+	n := 0
+	for i := range bits {
+		bits[i] = 15 // ^uint64(0)
+		n += 4
+		if n > rbf.ArrayMaxSize {
+			break
+		}
+	}
+
+	rb.Put(0, roaring.NewContainerBitmap(n, bits))
+	crun := roaring.NewContainerRun([]roaring.Interval16{roaring.Interval16{Start: 0, Last: 1<<16 - 1}})
+	rb.Put(1, crun)
+	rb.Put(2, roaring.NewContainerArray([]uint16{1, 1024, 1<<16 - 1}))
+
+	// setup to delete random bits down
+	values := rb.Slice()
+	valmap := make(map[uint64]bool)
+	for _, v := range values {
+		valmap[v] = true
+	}
+
+	_, err := tx.AddRoaring("x", rb)
+	if err != nil {
+		t.Errorf("Add Roaring Failed %v", err)
+	}
+
+	c, err := tx.Cursor("x")
+	if err != nil {
+		t.Fatal(err)
+	}
+	c.DebugSlowCheckAllPages()
+
+	if err := c.First(); err != nil {
+		t.Fatal(err)
+	}
+
+	exists, err := c.Contains(0x3)
+	if err != nil {
+		t.Fatalf("ERR:%v", err)
+	}
+	if !exists {
+		t.Fatalf("Should Contain %v", 0x3)
+	}
+
+	exists, err = c.Contains(0x4)
+	if err != nil {
+		t.Fatalf("ERR:%v", err)
+	}
+	if exists {
+		t.Fatalf("Should Not Contain %v", 0x4)
+	}
+
+	c.DebugSlowCheckAllPages()
+
+	_ = tx.ForEach("x", func(i uint64) error {
+		delete(valmap, i)
+		return nil
+	})
+
+	// check it is empty
+	if len(valmap) != 0 {
+		t.Fatalf("expected empty container, but see %v values left: '%#v'", len(valmap), valmap)
+	}
+}

rbf/rbf.go

@@ -45,10 +45,10 @@ const (
 
 	// ArrayMaxSize represents the maximum size of array containers.
 	// This is sligtly less than roaring to accommodate the page header.
-	ArrayMaxSize = 4080
+	ArrayMaxSize = 4079
 
 	// RLEMaxSize represents the maximum size of run length encoded containers.
-	RLEMaxSize = 2040
+	RLEMaxSize = 2039
 )
 
 // Page types.
@@ -98,7 +98,7 @@ func (typ ContainerType) String() string {
 const (
 	rootRecordPageHeaderSize = 12
 	rootRecordHeaderSize     = 4 + 2     // pgno, len(name)
-	leafCellHeaderSize       = 8 + 4 + 4 // key, type, count
+	leafCellHeaderSize       = 8 + 4 + 6 // key, type, count
 	branchCellSize           = 8 + 4 + 4 // key, flags, pgno
 )
 
@@ -481,15 +481,15 @@ func readLeafCell(page []byte, i int) leafCell {
 	cell.Key = *(*uint64)(unsafe.Pointer(&buf[0]))
 	cell.Type = ContainerType(*(*uint32)(unsafe.Pointer(&buf[8])))
 	cell.ElemN = int(*(*uint16)(unsafe.Pointer(&buf[12])))
-	cell.BitN = int(*(*uint16)(unsafe.Pointer(&buf[14])))
+	cell.BitN = int(*(*uint32)(unsafe.Pointer(&buf[14])))
 
 	switch cell.Type {
 	case ContainerTypeArray:
-		cell.Data = buf[16 : 16+(cell.ElemN*2)]
+		cell.Data = buf[18 : 18+(cell.ElemN*2)]
 	case ContainerTypeRLE:
-		cell.Data = buf[16 : 16+(cell.ElemN*4)]
+		cell.Data = buf[18 : 18+(cell.ElemN*4)]
 	case ContainerTypeBitmapPtr:
-		cell.Data = buf[16 : 16+4]
+		cell.Data = buf[18 : 18+4]
 	default:
 	}
 
@@ -519,9 +519,9 @@ func writeLeafCell(page []byte, i, offset int, cell leafCell) {
 	*(*uint64)(unsafe.Pointer(&page[offset])) = cell.Key
 	*(*uint32)(unsafe.Pointer(&page[offset+8])) = uint32(cell.Type)
 	*(*uint16)(unsafe.Pointer(&page[offset+12])) = uint16(cell.ElemN)
-	*(*uint16)(unsafe.Pointer(&page[offset+14])) = uint16(cell.BitN)
-	assert(offset+16+len(cell.Data) <= PageSize) // leaf cell write extends beyond page
-	copy(page[offset+16:], cell.Data)
+	*(*uint32)(unsafe.Pointer(&page[offset+14])) = uint32(cell.BitN)
+	assert(offset+18+len(cell.Data) <= PageSize) // leaf cell write extends beyond page
+	copy(page[offset+18:], cell.Data)
 }
 
 // branchCell represents a branch cell.

rbf/tx.go

@@ -1140,7 +1140,28 @@ func (tx *Tx) ForEachRange(name string, start, end uint64, fn func(uint64) error
 		case ContainerTypeBitmap:
 			for i, bits := range toArray64(cell.Data) {
 				for j := uint(0); j < 64; j++ {
-					if bits&(1<<j) != 0 {
+					if bits&(1<<j) == 0 {
+						continue
+					}
+
+					v := cell.Key<<16 | (uint64(i) * 64) | uint64(j)
+					if v < start {
+						continue
+					} else if v > end {
+						return nil
+					} else if err := fn(v); err != nil {
+						return err
+					}
+				}
+			}
+		case ContainerTypeBitmapPtr:
+			_, bm, err := c.tx.leafCellBitmap(toPgno(cell.Data))
+			if err != nil {
+				return err
+			}
+			for i, bits := range bm {
+				for j := uint(0); j < 64; j++ {
+					if bits&(1<<j) == 0 {
 						continue
 					}