Support looking up pools by frame and flagging bitmap entries as

reserved/free

kernel/mem/pmm/allocator/bitmap_allocator.go

@@ -22,6 +22,13 @@ var (
 	mapFn           = vmm.Map
 )
 
+type markAs bool
+
+const (
+	markReserved markAs = false
+	markFree            = true
+)
+
 type framePool struct {
 	// startFrame is the frame number for the first page in this pool.
 	// each free bitmap entry i corresponds to frame (startFrame + i).
@@ -145,6 +152,43 @@ func (alloc *BitmapAllocator) setupPoolBitmaps() *kernel.Error {
 	return nil
 }
 
+// markFrame updates the reservation flag for the bitmap entry that corresponds
+// to the supplied frame.
+func (alloc *BitmapAllocator) markFrame(poolIndex int, frame pmm.Frame, flag markAs) {
+	if poolIndex < 0 || frame > alloc.pools[poolIndex].endFrame {
+		return
+	}
+
+	// The offset in the block is given by: frame % 64. As the bitmap uses a
+	// big-ending representation we need to set the bit at index: 63 - offset
+	relFrame := frame - alloc.pools[poolIndex].startFrame
+	block := relFrame >> 6
+	mask := uint64(1 << (63 - (relFrame - block<<6)))
+	switch flag {
+	case markFree:
+		alloc.pools[poolIndex].freeBitmap[block] &^= mask
+		alloc.pools[poolIndex].freeCount++
+		alloc.reservedPages--
+	case markReserved:
+		alloc.pools[poolIndex].freeBitmap[block] |= mask
+		alloc.pools[poolIndex].freeCount--
+		alloc.reservedPages++
+	}
+}
+
+// poolForFrame returns the index of the pool that contains frame or -1 if
+// the frame is not contained in any of the available memory pools (e.g it
+// points to a reserved memory region).
+func (alloc *BitmapAllocator) poolForFrame(frame pmm.Frame) int {
+	for poolIndex, pool := range alloc.pools {
+		if frame >= pool.startFrame && frame <= pool.endFrame {
+			return poolIndex
+		}
+	}
+
+	return -1
+}
+
 // earlyAllocFrame is a helper that delegates a frame allocation request to the
 // early allocator instance. This function is passed as an argument to
 // vmm.SetFrameAllocator instead of earlyAllocator.AllocFrame. The latter

kernel/mem/pmm/allocator/bitmap_allocator_test.go

@@ -130,6 +130,93 @@ func TestSetupPoolBitmapsErrors(t *testing.T) {
 	})
 }
 
+func TestBitmapAllocatorMarkFrame(t *testing.T) {
+	var alloc = BitmapAllocator{
+		pools: []framePool{
+			{
+				startFrame: pmm.Frame(0),
+				endFrame:   pmm.Frame(127),
+				freeCount:  128,
+				freeBitmap: make([]uint64, 2),
+			},
+		},
+		totalPages: 128,
+	}
+
+	lastFrame := pmm.Frame(alloc.totalPages)
+	for frame := pmm.Frame(0); frame < lastFrame; frame++ {
+		alloc.markFrame(0, frame, markReserved)
+
+		block := uint64(frame / 64)
+		blockOffset := uint64(frame % 64)
+		bitIndex := (63 - blockOffset)
+		bitMask := uint64(1 << bitIndex)
+
+		if alloc.pools[0].freeBitmap[block]&bitMask != bitMask {
+			t.Errorf("[frame %d] expected block[%d], bit %d to be set", frame, block, bitIndex)
+		}
+
+		alloc.markFrame(0, frame, markFree)
+
+		if alloc.pools[0].freeBitmap[block]&bitMask != 0 {
+			t.Errorf("[frame %d] expected block[%d], bit %d to be unset", frame, block, bitIndex)
+		}
+	}
+
+	// Calling markFrame with a frame not part of the pool should be a no-op
+	alloc.markFrame(0, pmm.Frame(0xbadf00d), markReserved)
+	for blockIndex, block := range alloc.pools[0].freeBitmap {
+		if block != 0 {
+			t.Errorf("expected all blocks to be set to 0; block %d is set to %d", blockIndex, block)
+		}
+	}
+
+	// Calling markFrame with a negative pool index should be a no-op
+	alloc.markFrame(-1, pmm.Frame(0), markReserved)
+	for blockIndex, block := range alloc.pools[0].freeBitmap {
+		if block != 0 {
+			t.Errorf("expected all blocks to be set to 0; block %d is set to %d", blockIndex, block)
+		}
+	}
+}
+
+func TestBitmapAllocatorPoolForFrame(t *testing.T) {
+	var alloc = BitmapAllocator{
+		pools: []framePool{
+			{
+				startFrame: pmm.Frame(0),
+				endFrame:   pmm.Frame(63),
+				freeCount:  64,
+				freeBitmap: make([]uint64, 1),
+			},
+			{
+				startFrame: pmm.Frame(128),
+				endFrame:   pmm.Frame(191),
+				freeCount:  64,
+				freeBitmap: make([]uint64, 1),
+			},
+		},
+		totalPages: 128,
+	}
+
+	specs := []struct {
+		frame    pmm.Frame
+		expIndex int
+	}{
+		{pmm.Frame(0), 0},
+		{pmm.Frame(63), 0},
+		{pmm.Frame(64), -1},
+		{pmm.Frame(128), 1},
+		{pmm.Frame(192), -1},
+	}
+
+	for specIndex, spec := range specs {
+		if got := alloc.poolForFrame(spec.frame); got != spec.expIndex {
+			t.Errorf("[spec %d] expected to get pool index %d; got %d", specIndex, spec.expIndex, got)
+		}
+	}
+}
+
 func TestAllocatorPackageInit(t *testing.T) {
 	defer func() {
 		mapFn = vmm.Map