arch-image-builder/builder/lib/area.py
BigfootACA 067ee2d341 add more comments
Signed-off-by: BigfootACA <bigfoot@classfun.cn>
2024-05-20 09:56:42 +08:00

240 lines
6.0 KiB
Python

from typing import Self
from builder.lib.utils import round_up, round_down, size_to_bytes
from builder.lib.serializable import SerializableDict, SerializableList
class Area(SerializableDict):
start: int = -1
end: int = -1
size: int = -1
def set(self, start: int = -1, end: int = -1, size: int = -1) -> Self:
self.start, self.end, self.size = start, end, size
return self
def to_tuple(self) -> tuple[int, int, int]:
return self.start, self.end, self.size
def to_dict(self) -> dict:
return {
"start": self.start,
"end": self.end,
"size": self.size,
}
def reset(self) -> Self:
"""
Remove all fields
"""
self.set(-1, -1, -1)
return self
def from_dict(self, o: dict) -> Self:
"""
Load all fields from config
"""
self.reset()
if "start" in o: self.start = size_to_bytes(o["start"])
if "offset" in o: self.start = size_to_bytes(o["offset"])
if "end" in o: self.end = size_to_bytes(o["end"])
if "size" in o: self.size = size_to_bytes(o["size"])
if "length" in o: self.size = size_to_bytes(o["length"])
return self
def is_area_in(self, area: Self) -> bool:
"""
Is another area full in this area
"""
self.fixup()
area.fixup()
return (
(self.start <= area.start <= self.end) and
(self.start <= area.end <= self.end) and
(area.size <= self.size)
)
def fixup(self) -> Self:
"""
Fill missing fields
"""
if self.start >= 0 and self.end >= 0 and self.start > self.end + 1:
raise ValueError("start large than end")
if 0 <= self.end < self.size and self.size >= 0:
raise ValueError("size large than end")
if self.start >= 0 and self.end >= 0 and self.size >= 0:
if self.size != self.end - self.start + 1:
raise ValueError("bad size")
elif self.start >= 0 and self.end >= 0: # need size
self.size = self.end - self.start + 1
elif self.start >= 0 and self.size >= 0: # need end
self.end = self.start + self.size - 1
elif self.end >= 0 and self.size >= 0: # need start
self.start = self.end - self.size + 1
else:
raise ValueError("missing value")
return self
def __init__(self, start: int = -1, end: int = -1, size: int = -1, area: Self = None):
"""
Initialize a area
"""
super().__init__()
if area: start, end, size = area.to_tuple()
self.start, self.end, self.size = start, end, size
def convert(start: int = -1, end: int = -1, size: int = -1, area: Area = None) -> Area:
return Area(start, end, size, area).fixup()
def to_tuple(start: int = -1, end: int = -1, size: int = -1, area: Area = None) -> tuple[int, int, int]:
return convert(start, end, size, area).to_tuple()
class Areas(list[Area], SerializableList):
def is_area_in(self, area: Area) -> bool:
"""
Is an area fully in this areas
"""
return any(pool.is_area_in(area) for pool in self)
def merge(self) -> Self:
"""
Merge all areas
"""
idx = 0
self.sort(key=lambda x: (x.start, x.end))
while len(self) > 0:
curr = self[idx]
if curr.size <= 0:
self.remove(curr)
continue
if idx > 0:
last = self[idx - 1]
if last.end + 1 >= curr.start:
# last end equals to this start
ent = Area(last.start, curr.end)
ent.fixup()
self.remove(last)
self.remove(curr)
self.insert(idx - 1, ent)
idx -= 1
idx += 1
if idx >= len(self): break
return self
def lookup(
self,
start: int = -1,
end: int = -1,
size: int = -1,
area: Area = None,
) -> Area | None:
"""
Lookup an area with fields
"""
start, end, size = to_tuple(start, end, size, area)
for area in self:
if not (area.start <= start <= area.end): continue
if not (area.start <= end <= area.end): continue
if size > area.size: continue
return area
return None
def align(self, align: int) -> Self:
"""
Align all fields to value
"""
self.sort(key=lambda x: (x.start, x.end))
for area in self:
start = round_up(area.start, align)
end = round_down(area.end + 1, align) - 1
size = end - start + 1
if start >= end or size < align:
self.remove(area)
else:
area.set(start, end, size)
self.merge()
return self
def add(
self,
start: int = -1,
end: int = -1,
size: int = -1,
area: Area = None
) -> Area | None:
"""
Add an area to this areas
"""
if area: start, end, size = area.to_tuple()
cnt = (start >= 0) + (end >= 0) + (size >= 0)
if cnt < 2: raise ValueError("missing value")
r = convert(start, end, size)
if r.size <= 0: return None
self.append(r)
return r
def splice(
self,
start: int = -1,
end: int = -1,
size: int = -1,
area: Area = None,
) -> bool:
"""
Remove a range from areas
"""
start, end, size = to_tuple(start, end, size, area)
if len(self) <= 0: return False
rs = min(area.start for area in self)
re = max(area.end for area in self)
if start < rs: start = rs
if end > re: end = re
start, end, size = to_tuple(start, end)
target = self.lookup(start, end, size)
if target is None: return False
self.remove(target)
self.add(target.start, start - 1)
self.add(end + 1, target.end)
self.merge()
return True
def find(
self,
start: int = -1,
end: int = -1,
size: int = -1,
area: Area = None,
biggest: bool = True,
) -> Area | None:
"""
Find matched area
"""
if area: start, end, size = area.to_tuple()
cnt = (start >= 0) + (end >= 0) + (size >= 0)
if cnt >= 2:
area = convert(start, end, size)
return area if self.is_area_in(area) else None
use = Areas()
for free in self:
if start >= 0 and not (free.start <= start <= free.end): continue
if end >= 0 and not (free.start <= end <= free.end): continue
if size >= 0 and size > free.size: continue
use.add(area=free)
if biggest: use.sort(key=lambda x: x.size)
if len(use) <= 0: return None
target = use[0]
if start >= 0: target.start, target.end = start, -1
if end >= 0: target.start, target.end = -1, end
if size >= 0: target.end, target.size = -1, size
return target.fixup()
def to_list(self) -> list:
return self
def from_list(self, o: list) -> Self:
self.clear()
for i in o: self.append(Area().from_dict(i))
return self