Using more Python-ish constructions in BIRD linked-list accessors

1 files changed, 30 insertions, 51 deletions

diff --git a/bird-gdb.py b/bird-gdb.py
index fe7827fe..5ed99259 100644
--- a/bird-gdb.py
+++ b/bird-gdb.py
@@ -1,3 +1,6 @@
+import itertools
+import functools
+
 class BIRDPrinter:
     def __init__(self, val):
         self.val = val
@@ -198,12 +201,20 @@ class BIRDList:
         if str(self.tail_node["prev"].address) == '0x0':
             raise Exception("List tail is NULL")
 
-    def walk(self, do):
+    def __iter__(self):
         cur = self.head
         while cur.dereference() != self.tail_node:
-            do(cur)
+            yield cur
             cur = cur.dereference()["next"]
 
+    def __len__(self):
+        return sum([1 for _ in self])
+
+    def __getitem__(self, key):
+        for item in itertools.islice(self, key):
+            return item
+
+        raise KeyError("Not enough elements in list")
 
 class BIRDListLength(gdb.Function):
     """Returns length of the list, as in
@@ -211,13 +222,8 @@ class BIRDListLength(gdb.Function):
     def __init__(self):
         super(BIRDListLength, self).__init__("list_length")
 
-    def count(self, _):
-        self.cnt += 1
-
     def invoke(self, l):
-        self.cnt = 0
-        BIRDList(l).walk(self.count)
-        return self.cnt
+        return len(BIRDList(l))
 
 BIRDListLength()
 
@@ -226,30 +232,11 @@ class BIRDListItem(gdb.Function):
     def __init__(self):
         super(BIRDListItem, self).__init__("list_item")
 
-    class BLException(Exception):
-        def __init__(self, node, msg):
-            Exception.__init__(self, msg)
-            self.node = node
-
-    def count(self, node):
-        if self.cnt == self.pos:
-            raise self.BLException(node, "Node found")
-
-        self.cnt += 1
-
     def invoke(self, l, n, t=None, item="n"):
-        self.cnt = 0
-        self.pos = n
-        bl = BIRDList(l)
-        try:
-            bl.walk(self.count)
-        except self.BLException as e:
-            if t is None:
-                return e.node
-            else:
-                return BIRD.skip_back(t, item, e.node)
-
-        raise Exception("List too short")
+        if t is None:
+            return BIRDList(l)[n]
+        else:
+            return BIRD.skip_back(t, item, BIRDList(l)[n])
 
 BIRDListItem()
 
@@ -340,14 +327,13 @@ class BIRDSlabResource(BIRDResource):
         self.val = val.cast(self.slabtype)
         self.info = None
 
-    def count_heads_item(self, item):
-        self.hcnt += 1
-        self.used += item.dereference().cast(self.slheadtype)["num_full"]
-
     def count_heads(self, which):
         self.hcnt = 0
         self.used = 0
-        BIRDList(self.val[which + "_heads"]).walk(self.count_heads_item)
+        for item in BIRDList(self.val[which + "_heads"]):
+            self.hcnt += 1
+            self.used += item.dereference().cast(self.slheadtype)["num_full"]
+
         self.info[which + "_heads"] = self.hcnt
         self.info[which + "_used"] = self.used
         return (self.hcnt, self.used)
@@ -388,14 +374,14 @@ class BIRDPoolResource(BIRDResource):
         self.resptrtype = gdb.lookup_type("struct resource").pointer()
         self.page_size = gdb.lookup_symbol("page_size")[0].value()
         self.val = val.cast(self.pooltype)
-        self.items = None
+        self.inside = BIRDList(self.val["inside"])
 
-    def parse_inside(self, val):
-        self.items.append(BIRDNewResource(val.cast(self.resptrtype).dereference()))
+    def __iter__(self):
+        for val in self.inside:
+            yield BIRDNewResource(val.cast(self.resptrtype).dereference())
 
-    def parse(self):
-        self.items = []
-        BIRDList(self.val["inside"]).walk(self.parse_inside)
+    def __len__(self):
+        return len(self.inside)
 
     def free_pages(self):
         if str(self.val['pages']) == '0x0':
@@ -404,9 +390,6 @@ class BIRDPoolResource(BIRDResource):
             return self.val['pages'].dereference()['free']
 
     def memsize(self):
-        if self.items is None:
-            self.parse()
-
         sum = BIRDResourceSize(0, self.pooltype.sizeof, self.free_pages() * self.page_size)
 #        for i in self.items:
 #            sum += i.memsize()
@@ -414,13 +397,10 @@ class BIRDPoolResource(BIRDResource):
         return sum
 
     def __str__(self):
-        if self.items is None:
-            self.parse()
-
 #        for i in self.items:
 #            print(i)
 
-        return f"Resource pool {self.val.address} \"{self.val['name'].string()}\" containing {len(self.items)} items and {self.free_pages()} free pages"
+        return f"Resource pool {self.val.address} \"{self.val['name'].string()}\" containing {len(self)} items and {self.free_pages()} free pages"
 
 BIRDResourceMap = {
         "mbl_memsize": BIRDMBResource,
@@ -446,14 +426,13 @@ class BIRDResourcePrinter(BIRDPrinter):
     def __init__(self, val):
         super(BIRDResourcePrinter, self).__init__(val)
         self.resource = BIRDNewResource(val)
-        self.resource.parse()
         self.resourcetype = gdb.lookup_type("struct resource")
 
         if type(self.resource) == BIRDPoolResource:
             self.children = self.pool_children
 
     def pool_children(self):
-        return iter([ ("\n", i.val.cast(self.resourcetype)) for i in self.resource.items ])
+        return iter([ ("\n", i.val.cast(self.resourcetype)) for i in self.resource ])
 
     def to_string(self):
         return f"[ {str(self.resource.memsize())} ] {str(self.resource)}"