python_ffi_test.py (14023B)
1 #!/usr/bin/env python3 2 """Python FFI tests for the Arboricx C ABI. 3 4 Tests both the native fast-path bundle loader and the Tricu kernel fallback. 5 """ 6 import ctypes 7 import os 8 import sys 9 import time 10 11 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) 12 ZIG_DIR = os.path.dirname(SCRIPT_DIR) 13 lib_path = os.environ.get( 14 "ARBORICX_LIB", 15 os.path.join(ZIG_DIR, "zig-out", "lib", "libarboricx.so"), 16 ) 17 lib = ctypes.CDLL(lib_path) 18 19 # --- Lifecycle --- 20 lib.arboricx_init.restype = ctypes.c_void_p 21 lib.arboricx_free.argtypes = [ctypes.c_void_p] 22 23 # --- Tree construction --- 24 lib.arb_leaf.argtypes = [ctypes.c_void_p] 25 lib.arb_leaf.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 26 lib.arb_leaf.restype = ctypes.c_int 27 lib.arb_stem.argtypes = [ctypes.c_void_p, ctypes.c_uint64] 28 lib.arb_stem.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 29 lib.arb_stem.restype = ctypes.c_int 30 lib.arb_fork.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64] 31 lib.arb_fork.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 32 lib.arb_fork.restype = ctypes.c_int 33 lib.arb_app.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64] 34 lib.arb_app.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 35 lib.arb_app.restype = ctypes.c_int 36 37 # --- Reduction --- 38 lib.arb_reduce.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64] 39 lib.arb_reduce.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 40 lib.arb_reduce.restype = ctypes.c_int 41 42 # --- Codecs --- 43 lib.arb_of_number.argtypes = [ctypes.c_void_p, ctypes.c_uint64] 44 lib.arb_of_number.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 45 lib.arb_of_number.restype = ctypes.c_int 46 lib.arb_of_string.argtypes = [ctypes.c_void_p, ctypes.c_char_p] 47 lib.arb_of_string.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 48 lib.arb_of_string.restype = ctypes.c_int 49 lib.arb_of_bytes.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t] 50 lib.arb_of_bytes.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 51 lib.arb_of_bytes.restype = ctypes.c_int 52 lib.arb_of_list.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] 53 lib.arb_of_list.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 54 lib.arb_of_list.restype = ctypes.c_int 55 lib.arb_to_number.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64)] 56 lib.arb_to_number.restype = ctypes.c_int 57 lib.arb_to_string.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_uint8)), ctypes.POINTER(ctypes.c_size_t)] 58 lib.arb_to_string.restype = ctypes.c_int 59 lib.arb_to_bool.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_int)] 60 lib.arb_to_bool.restype = ctypes.c_int 61 lib.arboricx_free_buf.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t] 62 63 # --- Result unwrapping --- 64 lib.arb_unwrap_result.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64)] 65 lib.arb_unwrap_result.restype = ctypes.c_int 66 lib.arb_unwrap_host_value.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64)] 67 lib.arb_unwrap_host_value.restype = ctypes.c_int 68 69 # --- Kernel --- 70 lib.arb_kernel_root.argtypes = [ctypes.c_void_p] 71 lib.arb_kernel_root.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 72 lib.arb_kernel_root.restype = ctypes.c_int 73 74 # --- Native bundle loading --- 75 lib.arb_load_bundle.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t, ctypes.c_char_p] 76 lib.arb_load_bundle.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 77 lib.arb_load_bundle.restype = ctypes.c_int 78 lib.arb_load_bundle_default.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t] 79 lib.arb_load_bundle_default.argtypes += [ctypes.POINTER(ctypes.c_uint64)] 80 lib.arb_load_bundle_default.restype = ctypes.c_int 81 82 83 lib.arb_release.argtypes = [ctypes.c_void_p, ctypes.c_uint64] 84 lib.arb_release.restype = ctypes.c_int 85 lib.arb_collect.argtypes = [ctypes.c_void_p] 86 lib.arb_collect.restype = ctypes.c_int 87 lib.arb_abi_version.restype = ctypes.c_uint32 88 assert lib.arb_abi_version() == 2 89 90 def owned(name, *args): 91 out = ctypes.c_uint64() 92 status = getattr(lib, name)(*args, ctypes.byref(out)) 93 if status != 0: 94 raise RuntimeError(f"{name}: status {status}") 95 return out.value 96 97 ctx = lib.arboricx_init() 98 print("ctx init ok") 99 100 fixtures = os.path.join(ZIG_DIR, "..", "..", "test", "fixtures") 101 102 103 def read_bundle(name): 104 path = os.path.join(fixtures, name) 105 with open(path, "rb") as f: 106 return f.read() 107 108 109 def c_bytes(py_bytes): 110 arr = (ctypes.c_uint8 * len(py_bytes))(*py_bytes) 111 return arr 112 113 114 def to_string(ctx, root): 115 ptr = ctypes.POINTER(ctypes.c_uint8)() 116 length = ctypes.c_size_t() 117 if lib.arb_to_string(ctx, root, ctypes.byref(ptr), ctypes.byref(length)) != 0: 118 raise RuntimeError("to_string failed") 119 result = bytes(ptr[i] for i in range(length.value)) 120 lib.arboricx_free_buf(ctx, ptr, length.value) 121 return result.decode("utf-8") 122 123 124 def to_number(ctx, root): 125 out = ctypes.c_uint64() 126 if lib.arb_to_number(ctx, root, ctypes.byref(out)) != 0: 127 raise RuntimeError("to_number failed") 128 return out.value 129 130 131 def to_bool(ctx, root): 132 out = ctypes.c_int() 133 if lib.arb_to_bool(ctx, root, ctypes.byref(out)) != 0: 134 raise RuntimeError("to_bool failed") 135 return bool(out.value) 136 137 138 def kernel_run(bundle_bytes, args): 139 """Run via the Tricu kernel interpreter (slow, ~3s for append).""" 140 buf = c_bytes(bundle_bytes) 141 bundle_tree = owned("arb_of_bytes", ctx, buf, len(bundle_bytes)) 142 tag = owned("arb_of_number", ctx, 1) 143 arg_items = [] 144 for a in args: 145 arg_items.append(owned("arb_of_string", ctx, a.encode("utf-8"))) 146 current = owned("arb_leaf", ctx) 147 for item in reversed(arg_items): 148 current = owned("arb_fork", ctx, item, current) 149 app0 = owned("arb_app", ctx, owned("arb_kernel_root", ctx), tag) 150 app1 = owned("arb_app", ctx, app0, bundle_tree) 151 app2 = owned("arb_app", ctx, app1, current) 152 result = owned("arb_reduce", ctx, app2, 1_000_000_000) 153 ok = ctypes.c_int() 154 value = ctypes.c_uint64() 155 rest = ctypes.c_uint64() 156 if lib.arb_unwrap_result(ctx, result, ctypes.byref(ok), ctypes.byref(value), ctypes.byref(rest)) != 0: 157 raise RuntimeError("unwrap_result failed") 158 tag_num = ctypes.c_uint64() 159 payload = ctypes.c_uint64() 160 if lib.arb_unwrap_host_value(ctx, value.value, ctypes.byref(tag_num), ctypes.byref(payload)) != 0: 161 raise RuntimeError("unwrap_host_value failed") 162 return to_string(ctx, payload.value) 163 164 165 def native_run_default(bundle_bytes, args): 166 """Run via native bundle loader (fast, ~0.01s).""" 167 buf = c_bytes(bundle_bytes) 168 term = owned("arb_load_bundle_default", ctx, buf, len(bundle_bytes)) 169 if term == 0: 170 raise RuntimeError("load_bundle_default failed") 171 current = term 172 for a in args: 173 arg_tree = owned("arb_of_string", ctx, a.encode("utf-8")) 174 current = owned("arb_app", ctx, current, arg_tree) 175 result = owned("arb_reduce", ctx, current, 1_000_000_000) 176 return to_string(ctx, result) 177 178 179 def native_run_named(bundle_bytes, name, args): 180 """Run via native bundle loader with named export (fast).""" 181 buf = c_bytes(bundle_bytes) 182 term = owned("arb_load_bundle", ctx, buf, len(bundle_bytes), name.encode("utf-8")) 183 if term == 0: 184 raise RuntimeError(f"load_bundle({name!r}) failed") 185 current = term 186 for a in args: 187 arg_tree = owned("arb_of_string", ctx, a.encode("utf-8")) 188 current = owned("arb_app", ctx, current, arg_tree) 189 result = owned("arb_reduce", ctx, current, 1_000_000_000) 190 return to_string(ctx, result) 191 192 193 # ============================================================================ 194 # Tests 195 # ============================================================================ 196 197 all_ok = True 198 199 200 def check(label, got, want): 201 global all_ok 202 if got != want: 203 print(f"FAIL {label}: got {got!r}, want {want!r}") 204 all_ok = False 205 else: 206 print(f"PASS {label}: {got!r}") 207 208 209 # Test 1: id via kernel 210 print("\n--- Test 1: id (kernel path) ---") 211 bundle = read_bundle("id.arboricx") 212 t0 = time.time() 213 result = kernel_run(bundle, ["hello"]) 214 t1 = time.time() 215 check("id kernel", result, "hello") 216 print(f" time: {(t1 - t0) * 1000:.1f} ms") 217 218 # Test 2: id via native 219 print("\n--- Test 2: id (native path) ---") 220 t0 = time.time() 221 result = native_run_default(bundle, ["hello"]) 222 t1 = time.time() 223 check("id native", result, "hello") 224 print(f" time: {(t1 - t0) * 1000:.1f} ms") 225 226 # Test 3: append via kernel 227 print("\n--- Test 3: append (kernel path) ---") 228 bundle = read_bundle("append.arboricx") 229 t0 = time.time() 230 result = kernel_run(bundle, ["Hello, ", "world!"]) 231 t1 = time.time() 232 check("append kernel", result, "Hello, world!") 233 print(f" time: {(t1 - t0) * 1000:.1f} ms") 234 235 # Test 4: append via native 236 print("\n--- Test 4: append (native path) ---") 237 t0 = time.time() 238 result = native_run_default(bundle, ["Hello, ", "world!"]) 239 t1 = time.time() 240 check("append native", result, "Hello, world!") 241 print(f" time: {(t1 - t0) * 1000:.1f} ms") 242 243 # Test 5: append via native named export 244 print("\n--- Test 5: append via named export 'root' ---") 245 t0 = time.time() 246 result = native_run_named(bundle, "append", ["Hello, ", "world!"]) 247 t1 = time.time() 248 check("append named", result, "Hello, world!") 249 print(f" time: {(t1 - t0) * 1000:.1f} ms") 250 251 # Test 6: true / false via native 252 print("\n--- Test 6: true / false (native path) ---") 253 for name, expected in [("true.arboricx", True), ("false.arboricx", False)]: 254 bundle = read_bundle(name) 255 buf = c_bytes(bundle) 256 term = owned("arb_load_bundle_default", ctx, buf, len(bundle)) 257 result = owned("arb_reduce", ctx, term, 1_000_000_000) 258 check(f"{name} bool", to_bool(ctx, result), expected) 259 260 # Test 7: number roundtrip 261 print("\n--- Test 7: number roundtrip ---") 262 num_tree = owned("arb_of_number", ctx, 42) 263 check("number 42", to_number(ctx, num_tree), 42) 264 265 # Test 8: string roundtrip 266 print("\n--- Test 8: string roundtrip ---") 267 str_tree = owned("arb_of_string", ctx, b"hello") 268 check("string hello", to_string(ctx, str_tree), "hello") 269 270 # Independent canonical writer oracle, exercised through the public IO ABI. 271 # Build runtime graphs separately so shared/unshared equal children cannot be 272 # hidden by a matching bug in the native loader and serializer. 273 import hashlib 274 import random 275 import struct 276 import tempfile 277 from pathlib import Path 278 sys.path.insert(0, os.path.join(ZIG_DIR, "..", "..", "test")) 279 from executable_object import indexed_manifest, verify_payload, INDEXED 280 281 class Perms(ctypes.Structure): 282 _fields_ = [("read", ctypes.c_int), ("write", ctypes.c_int)] 283 lib.arb_run_io.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(Perms), ctypes.POINTER(ctypes.c_uint64)] 284 lib.arb_run_io.restype = ctypes.c_int 285 286 def pair(a, b): 287 return owned("arb_fork", ctx, a, b) 288 289 def action(tag, payload): 290 program = pair(owned("arb_of_string", ctx, b"tricuIO"), 291 pair(owned("arb_of_number", ctx, 1), 292 pair(owned("arb_of_number", ctx, tag), payload))) 293 result = owned("arb_run_io", ctx, program, ctypes.byref(Perms(1, 1))) 294 ok, value, rest = ctypes.c_int(), ctypes.c_uint64(), ctypes.c_uint64() 295 assert lib.arb_unwrap_result(ctx, result, ctypes.byref(ok), ctypes.byref(value), ctypes.byref(rest)) == 0 296 assert ok.value, to_string(ctx, value.value) 297 return value.value 298 299 rng = random.Random(510052) 300 with tempfile.TemporaryDirectory() as directory: 301 for case in range(24): 302 model, handles = [(0,)], [owned("arb_leaf", ctx)] 303 for i in range(1, 40): 304 tag = rng.randrange(3) 305 children = tuple(rng.randrange(i) for _ in range(tag)) 306 model.append((tag, *children)) 307 handles.append(owned(("arb_leaf", "arb_stem", "arb_fork")[tag], ctx, 308 *(handles[c] for c in children))) 309 canonical, by_key, memo = [], {}, {} 310 def visit(i): 311 if i in memo: 312 return memo[i] 313 key = (model[i][0], *(visit(c) for c in model[i][1:])) 314 if key not in by_key: 315 by_key[key] = len(canonical) 316 canonical.append(key) 317 memo[i] = by_key[key] 318 return memo[i] 319 root = visit(len(model)-1) 320 manifest = indexed_manifest(root) 321 nodes = struct.pack('>Q', len(canonical)) 322 for tag, *children in canonical: 323 nodes += struct.pack('>IB', 1 + 4*tag, tag) + b''.join(struct.pack('>I', c) for c in children) 324 expected = b'ARBORICX' + struct.pack('>HHIQQ', 1, 0, 2, 0, 32) 325 offset = 96 326 for kind, data in [(1, manifest), (2, nodes)]: 327 expected += struct.pack('>IHHHHQQI', kind, 1, 1, 0, 0, offset, len(data), 0) 328 offset += len(data) 329 expected += manifest + nodes 330 verify_payload(INDEXED, expected) 331 digest = hashlib.sha256(INDEXED.encode() + b'\0' + expected).hexdigest() 332 check(f"indexed independent hash {case}", to_string(ctx, action(51, handles[-1])), digest) 333 path = Path(directory) / str(case) 334 path_tree = owned("arb_of_string", ctx, str(path).encode()) 335 hash_tree = owned("arb_of_string", ctx, digest.encode()) 336 write = pair(path_tree, pair(hash_tree, handles[-1])) 337 assert to_bool(ctx, action(52, write)) is True 338 assert path.read_bytes() == expected 339 assert to_bool(ctx, action(52, write)) is False 340 decoded = action(50, pair(path_tree, hash_tree)) 341 assert to_string(ctx, action(51, decoded)) == digest 342 # Malformed bundle failure must not publish or clobber the output handle. 343 bad = expected[:31] 344 out = ctypes.c_uint64(123456) 345 assert lib.arb_load_bundle_default(ctx, c_bytes(bad), len(bad), ctypes.byref(out)) == 5 346 assert out.value == 123456 347 348 lib.arboricx_free(ctx) 349 350 if all_ok: 351 print("\nAll tests passed!") 352 sys.exit(0) 353 else: 354 print("\nSome tests failed!") 355 sys.exit(1)