Move RBTree to the rb_root extension object for better integration

This commit is contained in:
Gustavo Moreira
2024-08-15 07:55:45 -07:00
parent 87fe3444eb
commit 4513806a7e
2 changed files with 30 additions and 34 deletions
@@ -46,6 +46,7 @@ class LinuxKernelIntermedSymbols(intermed.IntermediateSymbolTable):
# Might not exist in older kernels or the current symbols
self.optional_set_type_class("mount", extensions.mount)
self.optional_set_type_class("mnt_namespace", extensions.mnt_namespace)
self.optional_set_type_class("rb_root", extensions.rb_root)
# Network
self.set_type_class("net", extensions.net)
@@ -425,36 +426,3 @@ class LinuxUtilities(interfaces.configuration.VersionableInterface):
kernel = context.modules[kernel_module_name]
return kernel
class RBTree(object):
"""Simple Red-Black tree abstraction"""
def __init__(self, root):
self.root = root
def _walk_nodes(self, root_node) -> Iterator[int]:
"""Traverses the Red-Black tree from the root node and yields a pointer to each
node in this tree.
Args:
root_node: A Red-Black tree node from which to start descending
Yields:
A pointer to every node descending from the specified root node
"""
if not root_node:
return
yield root_node
yield from self._walk_nodes(root_node.rb_left)
yield from self._walk_nodes(root_node.rb_right)
def get_nodes(self) -> Iterator[int]:
"""Yields a pointer to each node in the Red-Black tree
Yields:
A pointer to every node in the Red-Black tree
"""
yield from self._walk_nodes(root_node=self.root.rb_node)
@@ -1326,7 +1326,7 @@ class mnt_namespace(objects.StructType):
vmlinux = linux.LinuxUtilities.get_module_from_volobj_type(
self._context, self
)
for node in linux.RBTree(self.mounts).get_nodes():
for node in self.mounts.get_nodes():
mnt = linux.LinuxUtilities.container_of(
node, "mount", "mnt_list", vmlinux
)
@@ -1925,3 +1925,31 @@ class inode(objects.StructType):
The inode's file mode string
"""
return stat.filemode(self.i_mode)
class rb_root(objects.StructType):
def _walk_nodes(self, root_node) -> Iterator[int]:
"""Traverses the Red-Black tree from the root node and yields a pointer to each
node in this tree.
Args:
root_node: A Red-Black tree node from which to start descending
Yields:
A pointer to every node descending from the specified root node
"""
if not root_node:
return
yield root_node
yield from self._walk_nodes(root_node.rb_left)
yield from self._walk_nodes(root_node.rb_right)
def get_nodes(self) -> Iterator[int]:
"""Yields a pointer to each node in the Red-Black tree
Yields:
A pointer to every node in the Red-Black tree
"""
yield from self._walk_nodes(root_node=self.rb_node)