beam: harden rollback marker validation and add shape/type cells

wizard-work/beam/dsu_rollback.py

@@ -61,6 +61,8 @@ class DSU:
         return len(self.history)
 
     def rollback(self, marker: int) -> None:
+        if isinstance(marker, bool) or not isinstance(marker, int):
+            raise TypeError("marker must be int")
         if marker < 0 or marker > len(self.history):
             raise ValueError("invalid marker")
         while len(self.history) > marker:

wizard-work/beam/test_dsu_rollback.py

@@ -91,5 +91,27 @@ class TestDSURollback(unittest.TestCase):
             DSU(-1)
 
 
+    def test_shape_union_by_size(self):
+        d = DSU(8)
+        d.union(0,1); d.union(2,3); d.union(4,5); d.union(6,7)
+        d.union(1,2); d.union(4,0)
+        # union by size: attach smaller root under larger -> parent[4]==0, depth 2
+        self.assertEqual(d.parent[4], 0)
+        # deepest depth
+        maxdepth = 0
+        for i in range(8):
+            depth=0; cur=i
+            while d.parent[cur]!=cur:
+                cur=d.parent[cur]; depth+=1
+            maxdepth=max(maxdepth,depth)
+        self.assertEqual(maxdepth, 2)
+
+    def test_rollback_marker_type(self):
+        d = DSU(2)
+        for bad in (True, False, 1.5, "x"):
+            with self.assertRaises((TypeError, ValueError)):
+                d.rollback(bad)
+
+
 if __name__ == "__main__":
     unittest.main()