1# SPDX-License-Identifier: GPL-2.0
2#
3# Copyright 2019 Google LLC.
4
5import gdb
6
7from linux import utils
8
9rb_root_type = utils.CachedType("struct rb_root")
10rb_node_type = utils.CachedType("struct rb_node")
11
12def rb_inorder_for_each(root):
13    def inorder(node):
14        if node:
15            yield from inorder(node['rb_left'])
16            yield node
17            yield from inorder(node['rb_right'])
18
19    yield from inorder(root['rb_node'])
20
21def rb_inorder_for_each_entry(root, gdbtype, member):
22    for node in rb_inorder_for_each(root):
23        yield utils.container_of(node, gdbtype, member)
24
25def rb_first(root):
26    if root.type == rb_root_type.get_type():
27        node = root.address.cast(rb_root_type.get_type().pointer())
28    elif root.type != rb_root_type.get_type().pointer():
29        raise gdb.GdbError("Must be struct rb_root not {}".format(root.type))
30
31    node = root['rb_node']
32    if node == 0:
33        return None
34
35    while node['rb_left']:
36        node = node['rb_left']
37
38    return node
39
40
41def rb_last(root):
42    if root.type == rb_root_type.get_type():
43        node = root.address.cast(rb_root_type.get_type().pointer())
44    elif root.type != rb_root_type.get_type().pointer():
45        raise gdb.GdbError("Must be struct rb_root not {}".format(root.type))
46
47    node = root['rb_node']
48    if node == 0:
49        return None
50
51    while node['rb_right']:
52        node = node['rb_right']
53
54    return node
55
56
57def rb_parent(node):
58    parent = gdb.Value(node['__rb_parent_color'] & ~3)
59    return parent.cast(rb_node_type.get_type().pointer())
60
61
62def rb_empty_node(node):
63    return node['__rb_parent_color'] == node.address
64
65
66def rb_next(node):
67    if node.type == rb_node_type.get_type():
68        node = node.address.cast(rb_node_type.get_type().pointer())
69    elif node.type != rb_node_type.get_type().pointer():
70        raise gdb.GdbError("Must be struct rb_node not {}".format(node.type))
71
72    if rb_empty_node(node):
73        return None
74
75    if node['rb_right']:
76        node = node['rb_right']
77        while node['rb_left']:
78            node = node['rb_left']
79        return node
80
81    parent = rb_parent(node)
82    while parent and node == parent['rb_right']:
83            node = parent
84            parent = rb_parent(node)
85
86    return parent
87
88
89def rb_prev(node):
90    if node.type == rb_node_type.get_type():
91        node = node.address.cast(rb_node_type.get_type().pointer())
92    elif node.type != rb_node_type.get_type().pointer():
93        raise gdb.GdbError("Must be struct rb_node not {}".format(node.type))
94
95    if rb_empty_node(node):
96        return None
97
98    if node['rb_left']:
99        node = node['rb_left']
100        while node['rb_right']:
101            node = node['rb_right']
102        return node.dereference()
103
104    parent = rb_parent(node)
105    while parent and node == parent['rb_left'].dereference():
106            node = parent
107            parent = rb_parent(node)
108
109    return parent
110
111
112class LxRbFirst(gdb.Function):
113    """Lookup and return a node from an RBTree
114
115$lx_rb_first(root): Return the node at the given index.
116If index is omitted, the root node is dereferenced and returned."""
117
118    def __init__(self):
119        super(LxRbFirst, self).__init__("lx_rb_first")
120
121    def invoke(self, root):
122        result = rb_first(root)
123        if result is None:
124            raise gdb.GdbError("No entry in tree")
125
126        return result
127
128
129LxRbFirst()
130
131
132class LxRbLast(gdb.Function):
133    """Lookup and return a node from an RBTree.
134
135$lx_rb_last(root): Return the node at the given index.
136If index is omitted, the root node is dereferenced and returned."""
137
138    def __init__(self):
139        super(LxRbLast, self).__init__("lx_rb_last")
140
141    def invoke(self, root):
142        result = rb_last(root)
143        if result is None:
144            raise gdb.GdbError("No entry in tree")
145
146        return result
147
148
149LxRbLast()
150
151
152class LxRbNext(gdb.Function):
153    """Lookup and return a node from an RBTree.
154
155$lx_rb_next(node): Return the node at the given index.
156If index is omitted, the root node is dereferenced and returned."""
157
158    def __init__(self):
159        super(LxRbNext, self).__init__("lx_rb_next")
160
161    def invoke(self, node):
162        result = rb_next(node)
163        if result is None:
164            raise gdb.GdbError("No entry in tree")
165
166        return result
167
168
169LxRbNext()
170
171
172class LxRbPrev(gdb.Function):
173    """Lookup and return a node from an RBTree.
174
175$lx_rb_prev(node): Return the node at the given index.
176If index is omitted, the root node is dereferenced and returned."""
177
178    def __init__(self):
179        super(LxRbPrev, self).__init__("lx_rb_prev")
180
181    def invoke(self, node):
182        result = rb_prev(node)
183        if result is None:
184            raise gdb.GdbError("No entry in tree")
185
186        return result
187
188
189LxRbPrev()
190