python_ffi_test.py (15584B)
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() == 3 89 assert not hasattr(lib, "arb_run_io"), "retired v1 entrypoint remains exported" 90 91 def owned(name, *args): 92 out = ctypes.c_uint64() 93 status = getattr(lib, name)(*args, ctypes.byref(out)) 94 if status != 0: 95 raise RuntimeError(f"{name}: status {status}") 96 return out.value 97 98 ctx = lib.arboricx_init() 99 print("ctx init ok") 100 101 fixtures = os.path.join(ZIG_DIR, "..", "..", "test", "fixtures") 102 103 104 def read_bundle(name): 105 path = os.path.join(fixtures, name) 106 with open(path, "rb") as f: 107 return f.read() 108 109 110 def c_bytes(py_bytes): 111 arr = (ctypes.c_uint8 * len(py_bytes))(*py_bytes) 112 return arr 113 114 115 def to_string(ctx, root): 116 ptr = ctypes.POINTER(ctypes.c_uint8)() 117 length = ctypes.c_size_t() 118 if lib.arb_to_string(ctx, root, ctypes.byref(ptr), ctypes.byref(length)) != 0: 119 raise RuntimeError("to_string failed") 120 result = bytes(ptr[i] for i in range(length.value)) 121 lib.arboricx_free_buf(ctx, ptr, length.value) 122 return result.decode("utf-8") 123 124 125 def to_number(ctx, root): 126 out = ctypes.c_uint64() 127 if lib.arb_to_number(ctx, root, ctypes.byref(out)) != 0: 128 raise RuntimeError("to_number failed") 129 return out.value 130 131 132 def to_bool(ctx, root): 133 out = ctypes.c_int() 134 if lib.arb_to_bool(ctx, root, ctypes.byref(out)) != 0: 135 raise RuntimeError("to_bool failed") 136 return bool(out.value) 137 138 139 def kernel_run(bundle_bytes, args): 140 """Run via the Tricu kernel interpreter (slow, ~3s for append).""" 141 buf = c_bytes(bundle_bytes) 142 bundle_tree = owned("arb_of_bytes", ctx, buf, len(bundle_bytes)) 143 tag = owned("arb_of_number", ctx, 1) 144 arg_items = [] 145 for a in args: 146 arg_items.append(owned("arb_of_string", ctx, a.encode("utf-8"))) 147 current = owned("arb_leaf", ctx) 148 for item in reversed(arg_items): 149 current = owned("arb_fork", ctx, item, current) 150 app0 = owned("arb_app", ctx, owned("arb_kernel_root", ctx), tag) 151 app1 = owned("arb_app", ctx, app0, bundle_tree) 152 app2 = owned("arb_app", ctx, app1, current) 153 result = owned("arb_reduce", ctx, app2, 1_000_000_000) 154 ok = ctypes.c_int() 155 value = ctypes.c_uint64() 156 rest = ctypes.c_uint64() 157 if lib.arb_unwrap_result(ctx, result, ctypes.byref(ok), ctypes.byref(value), ctypes.byref(rest)) != 0: 158 raise RuntimeError("unwrap_result failed") 159 tag_num = ctypes.c_uint64() 160 payload = ctypes.c_uint64() 161 if lib.arb_unwrap_host_value(ctx, value.value, ctypes.byref(tag_num), ctypes.byref(payload)) != 0: 162 raise RuntimeError("unwrap_host_value failed") 163 return to_string(ctx, payload.value) 164 165 166 def native_run_default(bundle_bytes, args): 167 """Run via native bundle loader (fast, ~0.01s).""" 168 buf = c_bytes(bundle_bytes) 169 term = owned("arb_load_bundle_default", ctx, buf, len(bundle_bytes)) 170 if term == 0: 171 raise RuntimeError("load_bundle_default failed") 172 current = term 173 for a in args: 174 arg_tree = owned("arb_of_string", ctx, a.encode("utf-8")) 175 current = owned("arb_app", ctx, current, arg_tree) 176 result = owned("arb_reduce", ctx, current, 1_000_000_000) 177 return to_string(ctx, result) 178 179 180 def native_run_named(bundle_bytes, name, args): 181 """Run via native bundle loader with named export (fast).""" 182 buf = c_bytes(bundle_bytes) 183 term = owned("arb_load_bundle", ctx, buf, len(bundle_bytes), name.encode("utf-8")) 184 if term == 0: 185 raise RuntimeError(f"load_bundle({name!r}) failed") 186 current = term 187 for a in args: 188 arg_tree = owned("arb_of_string", ctx, a.encode("utf-8")) 189 current = owned("arb_app", ctx, current, arg_tree) 190 result = owned("arb_reduce", ctx, current, 1_000_000_000) 191 return to_string(ctx, result) 192 193 194 # ============================================================================ 195 # Tests 196 # ============================================================================ 197 198 all_ok = True 199 200 201 def check(label, got, want): 202 global all_ok 203 if got != want: 204 print(f"FAIL {label}: got {got!r}, want {want!r}") 205 all_ok = False 206 else: 207 print(f"PASS {label}: {got!r}") 208 209 210 # Test 1: id via kernel 211 print("\n--- Test 1: id (kernel path) ---") 212 bundle = read_bundle("id.arboricx") 213 t0 = time.time() 214 result = kernel_run(bundle, ["hello"]) 215 t1 = time.time() 216 check("id kernel", result, "hello") 217 print(f" time: {(t1 - t0) * 1000:.1f} ms") 218 219 # Test 2: id via native 220 print("\n--- Test 2: id (native path) ---") 221 t0 = time.time() 222 result = native_run_default(bundle, ["hello"]) 223 t1 = time.time() 224 check("id native", result, "hello") 225 print(f" time: {(t1 - t0) * 1000:.1f} ms") 226 227 # Test 3: append via kernel 228 print("\n--- Test 3: append (kernel path) ---") 229 bundle = read_bundle("append.arboricx") 230 t0 = time.time() 231 result = kernel_run(bundle, ["Hello, ", "world!"]) 232 t1 = time.time() 233 check("append kernel", result, "Hello, world!") 234 print(f" time: {(t1 - t0) * 1000:.1f} ms") 235 236 # Test 4: append via native 237 print("\n--- Test 4: append (native path) ---") 238 t0 = time.time() 239 result = native_run_default(bundle, ["Hello, ", "world!"]) 240 t1 = time.time() 241 check("append native", result, "Hello, world!") 242 print(f" time: {(t1 - t0) * 1000:.1f} ms") 243 244 # Test 5: append via native named export 245 print("\n--- Test 5: append via named export 'root' ---") 246 t0 = time.time() 247 result = native_run_named(bundle, "append", ["Hello, ", "world!"]) 248 t1 = time.time() 249 check("append named", result, "Hello, world!") 250 print(f" time: {(t1 - t0) * 1000:.1f} ms") 251 252 # Test 6: true / false via native 253 print("\n--- Test 6: true / false (native path) ---") 254 for name, expected in [("true.arboricx", True), ("false.arboricx", False)]: 255 bundle = read_bundle(name) 256 buf = c_bytes(bundle) 257 term = owned("arb_load_bundle_default", ctx, buf, len(bundle)) 258 result = owned("arb_reduce", ctx, term, 1_000_000_000) 259 check(f"{name} bool", to_bool(ctx, result), expected) 260 261 # Test 7: number roundtrip 262 print("\n--- Test 7: number roundtrip ---") 263 num_tree = owned("arb_of_number", ctx, 42) 264 check("number 42", to_number(ctx, num_tree), 42) 265 266 # Test 8: string roundtrip 267 print("\n--- Test 8: string roundtrip ---") 268 str_tree = owned("arb_of_string", ctx, b"hello") 269 check("string hello", to_string(ctx, str_tree), "hello") 270 271 # Independent canonical writer oracle, exercised through the public IO ABI. 272 # Build runtime graphs separately so shared/unshared equal children cannot be 273 # hidden by a matching bug in the native loader and serializer. 274 import hashlib 275 import random 276 import struct 277 import tempfile 278 from pathlib import Path 279 sys.path.insert(0, os.path.join(ZIG_DIR, "..", "..", "test")) 280 from executable_object import indexed_manifest, verify_payload, INDEXED 281 282 lib.arb_run_io_v2.argtypes = [ctypes.c_void_p, ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint64)] 283 lib.arb_run_io_v2.restype = ctypes.c_int 284 lib.arb_get_fork_children.argtypes = [ctypes.c_void_p, ctypes.c_uint64, 285 ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64)] 286 lib.arb_get_fork_children.restype = ctypes.c_int 287 288 def pair(a, b): 289 return owned("arb_fork", ctx, a, b) 290 291 def fork_children(value): 292 a, b = ctypes.c_uint64(), ctypes.c_uint64() 293 assert lib.arb_get_fork_children(ctx, value, ctypes.byref(a), ctypes.byref(b)) == 0 294 return a.value, b.value 295 296 def execute(action): 297 program = pair(owned("arb_of_string", ctx, b"tricuIO"), 298 pair(owned("arb_of_number", ctx, 2), action)) 299 result = owned("arb_run_io_v2", ctx, program) 300 ok, value, rest = ctypes.c_int(), ctypes.c_uint64(), ctypes.c_uint64() 301 assert lib.arb_unwrap_result(ctx, result, ctypes.byref(ok), ctypes.byref(value), ctypes.byref(rest)) == 0 302 assert ok.value, "v2 operation failed" 303 return value.value 304 305 # Only the scoped storage composition comes from the authenticated compiler 306 # library. The canonical DAG and expected bytes/digest below remain independent. 307 fixture = (Path(ZIG_DIR).parent.parent / 'test/fixtures/compiler-library.arboricx').read_bytes() 308 def export(name): 309 return owned('arb_load_bundle', ctx, c_bytes(fixture), len(fixture), name.encode()) 310 publish_tree = export('io_v2_storage.publishTree') 311 read_tree = export('io_v2_storage.readTree') 312 def call(function, *arguments): 313 for argument in arguments: 314 function = owned('arb_app', ctx, function, argument) 315 return execute(function) 316 limit = owned('arb_of_number', ctx, 1048576) 317 def hash_tree(value): 318 # Independently encode indexed-tree revision 1 hash request. 319 request = pair(owned('arb_of_string', ctx, b'indexed-tree'), 320 pair(owned('arb_of_number', ctx, 1), pair(owned('arb_of_string', ctx, b'hash'), pair(limit, value)))) 321 info = execute(pair(owned('arb_of_number', ctx, 10), request)) 322 size, tail = fork_children(info) 323 digest, _ = fork_children(tail) 324 return to_number(ctx, size), to_string(ctx, digest) 325 326 rng = random.Random(510052) 327 with tempfile.TemporaryDirectory() as directory: 328 for case in range(24): 329 model, handles = [(0,)], [owned("arb_leaf", ctx)] 330 for i in range(1, 40): 331 tag = rng.randrange(3) 332 children = tuple(rng.randrange(i) for _ in range(tag)) 333 model.append((tag, *children)) 334 handles.append(owned(("arb_leaf", "arb_stem", "arb_fork")[tag], ctx, 335 *(handles[c] for c in children))) 336 canonical, by_key, memo = [], {}, {} 337 def visit(i): 338 if i in memo: 339 return memo[i] 340 key = (model[i][0], *(visit(c) for c in model[i][1:])) 341 if key not in by_key: 342 by_key[key] = len(canonical) 343 canonical.append(key) 344 memo[i] = by_key[key] 345 return memo[i] 346 root = visit(len(model)-1) 347 manifest = indexed_manifest(root) 348 nodes = struct.pack('>Q', len(canonical)) 349 for tag, *children in canonical: 350 nodes += struct.pack('>IB', 1 + 4*tag, tag) + b''.join(struct.pack('>I', c) for c in children) 351 expected = b'ARBORICX' + struct.pack('>HHIQQ', 1, 0, 2, 0, 32) 352 offset = 96 353 for kind, data in [(1, manifest), (2, nodes)]: 354 expected += struct.pack('>IHHHHQQI', kind, 1, 1, 0, 0, offset, len(data), 0) 355 offset += len(data) 356 expected += manifest + nodes 357 verify_payload(INDEXED, expected) 358 digest = hashlib.sha256(INDEXED.encode() + b'\0' + expected).hexdigest() 359 check(f"indexed independent hash {case}", hash_tree(handles[-1])[1], digest) 360 path = Path(directory) / str(case) 361 path_tree = owned("arb_of_string", ctx, str(path).encode()) 362 hash_tree_ = owned("arb_of_string", ctx, digest.encode()) 363 directory_tree = owned('arb_of_string', ctx, directory.encode()) 364 assert hash_tree(handles[-1])[0] == len(expected) 365 assert to_bool(ctx, call(publish_tree, directory_tree, path_tree, limit, hash_tree_, handles[-1])) is True 366 assert path.read_bytes() == expected 367 assert to_bool(ctx, call(publish_tree, directory_tree, path_tree, limit, hash_tree_, handles[-1])) is False 368 decoded = call(read_tree, directory_tree, path_tree, limit, hash_tree_) 369 assert hash_tree(decoded) == (len(expected), digest) 370 # Malformed bundle failure must not publish or clobber the output handle. 371 bad = expected[:31] 372 out = ctypes.c_uint64(123456) 373 assert lib.arb_load_bundle_default(ctx, c_bytes(bad), len(bad), ctypes.byref(out)) == 5 374 assert out.value == 123456 375 376 lib.arboricx_free(ctx) 377 378 if all_ok: 379 print("\nAll tests passed!") 380 sys.exit(0) 381 else: 382 print("\nSome tests failed!") 383 sys.exit(1)