|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import re |
| 4 | +from pathlib import Path |
| 5 | + |
| 6 | + |
| 7 | +ROOT = Path(__file__).resolve().parent.parent |
| 8 | +PATHS = ( |
| 9 | + "vortex-array/src/arrays/filter/execute", |
| 10 | + "vortex-array/src/arrays/fixed_width/filter.rs", |
| 11 | + "vortex-array/src/arrays/fixed_width/take", |
| 12 | + "vortex-array/src/arrays/interleave/execute", |
| 13 | + "vortex-array/src/patches.rs", |
| 14 | + "vortex-array/src/scalar_fn/fns", |
| 15 | +) |
| 16 | +TYPES = r"(?:ByteBuffer|Buffer|BitBuffer|ByteBufferMut|BufferMut|BitBufferMut)" |
| 17 | +METHODS = "|".join( |
| 18 | + ( |
| 19 | + "with_capacity", |
| 20 | + "with_capacity_aligned", |
| 21 | + "with_capacity_preferred_aligned", |
| 22 | + "zeroed", |
| 23 | + "zeroed_aligned", |
| 24 | + "empty", |
| 25 | + "empty_aligned", |
| 26 | + "copy_from", |
| 27 | + "copy_from_aligned", |
| 28 | + "full", |
| 29 | + "new_set", |
| 30 | + "new_unset", |
| 31 | + "collect_bool", |
| 32 | + "collect_bool_multiversioned", |
| 33 | + "from_trusted_len_iter", |
| 34 | + "try_from_trusted_len_iter", |
| 35 | + ) |
| 36 | +) |
| 37 | +STATIC_ALLOCATION = re.compile(rf"\b{TYPES}::(?:{METHODS})\s*\(") |
| 38 | + |
| 39 | + |
| 40 | +def rust_files(path: Path): |
| 41 | + files = path.rglob("*.rs") if path.is_dir() else (path,) |
| 42 | + return ( |
| 43 | + file |
| 44 | + for file in files |
| 45 | + if file.name != "tests.rs" and "tests" not in file.relative_to(ROOT).parts |
| 46 | + ) |
| 47 | + |
| 48 | + |
| 49 | +failures = [] |
| 50 | +for relative in PATHS: |
| 51 | + for file in rust_files(ROOT / relative): |
| 52 | + for number, line in enumerate(file.read_text().splitlines(), 1): |
| 53 | + if STATIC_ALLOCATION.search(line): |
| 54 | + failures.append(f"{file.relative_to(ROOT)}:{number}: {line.strip()}") |
| 55 | + |
| 56 | +if failures: |
| 57 | + print("Engine buffers must use an allocator-aware constructor ending in `_in`.") |
| 58 | + print("\n".join(failures)) |
| 59 | + raise SystemExit(1) |
0 commit comments