From 9b4ff5641a583ca4fa305165e460213d8201b2d9 Mon Sep 17 00:00:00 2001
From: Helmut Grohne <helmut@subdivi.de>
Date: Mon, 27 May 2024 17:01:59 +0200
Subject: add IDAllocation.reserve method

Allow reserving a particular range instead of allocating a suitable
large range of an IDAllocation. This is useful when a directory
hierarchy defines the allocation and we merely want to verify it to be
assigned.
---
 linuxnamespaces/__init__.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

(limited to 'linuxnamespaces')

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index b50f113..a39ab1d 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -136,6 +136,41 @@ class IDAllocation:
         """
         return IDMapping(target, self.allocate(count), count)
 
+    def reserve(self, start: int, count: int) -> None:
+        """Reserve (and remove) the given range from this allocation. If the
+        range is not fully contained in this allocation, a ValueError is
+        raised.
+        """
+        if count < 0:
+            raise ValueError("nagative count")
+        index = bisect.bisect_right(self.ranges, (start, float("inf"))) - 1
+        if index < 0:
+            raise ValueError("range to reserve not found")
+        cur_start, cur_count = self.ranges[index]
+        assert cur_start <= start
+        if cur_start == start:
+            # Requested range starts at range boundary
+            if cur_count < count:
+                raise ValueError("range to reserve not found")
+            if cur_count == count:
+                # Requested range matches a range exactly
+                del self.ranges[index]
+            else:
+                # Requested range is a head of the matched range
+                self.ranges[index] = (start + count, cur_count - count)
+        elif cur_start + cur_count >= start + count:
+            # Requested range fits into a matched range
+            self.ranges[index] = (cur_start, start - cur_start)
+            if cur_start + cur_count > start + count:
+                # Requested range punches a hole into a matched range
+                self.ranges.insert(
+                    index + 1,
+                    (start + count, cur_start + cur_count - (start + count))
+                )
+            # else: Requested range is a tail of a matched range
+        else:
+            raise ValueError("range to reserve not found")
+
 
 def newidmap(
     kind: typing.Literal["uid", "gid"],
-- 
cgit v1.2.3