Accept raw allocation sequences and streamline the public API

pyproject.toml

@@ -50,6 +50,7 @@ all = [
 dev = [
   "ipykernel>=6.30.0",
   "ipywidgets>=8.1.0",
+  "nanobind>=2.0",
   "nbconvert>=7.16.0",
   "pre-commit>=4.0.0",
   "pytest>=8.0",

src/python/omnimalloc/_allocate.py

@@ -3,37 +3,59 @@
 #
 
 
-from typing import TypeVar, cast
+from collections.abc import Sequence
+from typing import TypeAlias, TypeVar, cast, overload
 
 from .allocators import DEFAULT_ALLOCATOR, BaseAllocator
-from .primitives import Memory, Pool, System
+from .primitives import Allocation, Memory, Pool, System
 from .validate import validate_allocation
 
+AllocatorLike: TypeAlias = BaseAllocator | type[BaseAllocator] | str | None
+
 T = TypeVar("T", System, Memory, Pool)
 
 
+@overload
 def allocate(
     entity: T,
-    allocator: BaseAllocator | type[BaseAllocator] | str | None = None,
-    *,
+    allocator: AllocatorLike = None,
+    validate: bool = False,
+) -> T: ...
+
+
+@overload
+def allocate(
+    entity: Sequence[Allocation],
+    allocator: AllocatorLike = None,
     validate: bool = False,
-) -> T:
-    """Return the entity (System, Memory, or Pool) with offsets assigned.
+) -> tuple[Allocation, ...]: ...
+
 
-    `allocator` accepts an instance, a class, a registry name, or `None`
-    (the default allocator); `validate=True` additionally runs
-    `validate_allocation` on the result.
+def allocate(
+    entity: System | Memory | Pool | Sequence[Allocation],
+    allocator: AllocatorLike = None,
+    validate: bool = False,
+) -> System | Memory | Pool | tuple[Allocation, ...]:
+    """Return the entity with offsets assigned.
+
+    Accepts a System, Memory, or Pool (returned as the same type) or a raw
+    sequence of Allocations (returned as a tuple in input order). `allocator`
+    accepts an instance, a class, a registry name, or `None` (the default
+    allocator); `validate=True` additionally runs `validate_allocation` on
+    the result.
     """
 
     if allocator is None:
         allocator = DEFAULT_ALLOCATOR
 
-    resolved = BaseAllocator.resolve(allocator)
+    resolved = cast("BaseAllocator", BaseAllocator.resolve(allocator))
 
-    # ty doesn't understand that TypeVar T (System|Memory|Pool) all have allocate method
-    allocated = entity.allocate(resolved)  # type: ignore[invalid-argument-type]
+    if isinstance(entity, System | Memory | Pool):
+        allocated = entity.allocate(resolved)
+    else:
+        allocated = Pool.from_allocations(entity).allocate(resolved).allocations
 
     if validate:
         validate_allocation(allocated)
 
-    return cast("T", allocated)
+    return allocated

src/python/omnimalloc/allocators/genetic.py

@@ -44,7 +44,6 @@ class GeneticAllocator(GreedyAllocator):
 
     def __init__(
         self,
-        *,
         seed: int = DEFAULT_SEED,
         population_size: int = 100,
         max_generations: int = 50,

src/python/omnimalloc/allocators/greedy.py

@@ -77,7 +77,7 @@ class GreedyByAllAllocator(GreedyAllocator):
     `num_threads=None` uses all cores.
     """
 
-    def __init__(self, *, num_threads: int | None = None) -> None:
+    def __init__(self, num_threads: int | None = None) -> None:
         ensure_valid_num_threads(num_threads)
         self._num_threads = num_threads
 

src/python/omnimalloc/allocators/greedy_base.py

@@ -72,7 +72,6 @@ def _allocate(
 def allocate_parallel(
     allocations: tuple[Allocation, ...],
     variants: tuple[BaseAllocator, ...],
-    *,
     num_threads: int | None = None,
 ) -> tuple[Allocation, ...]:
     """Run each variant and return the lowest peak memory results.

src/python/omnimalloc/allocators/hillclimb.py

@@ -33,7 +33,6 @@ class HillClimbAllocator(GreedyAllocator):
 
     def __init__(
         self,
-        *,
         max_iterations: int = 100,
         seed: int = DEFAULT_SEED,
         acceptance_temperature: float = 2.0,

src/python/omnimalloc/allocators/minimalloc.py

@@ -51,7 +51,7 @@ class MinimallocAllocator(BaseAllocator):
     supports_vector_time = False
 
     def __init__(
-        self, *, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB
+        self, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB
     ) -> None:
         _require_minimalloc()
         ensure_valid_timeout(timeout)

src/python/omnimalloc/allocators/omni.py

@@ -23,7 +23,7 @@ class OmniAllocator(BaseAllocator):
 
     supports_vector_time = True
 
-    def __init__(self, *, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None:
+    def __init__(self, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None:
         ensure_valid_budget(linearize_budget, name="linearize_budget")
         self._linearize_budget = linearize_budget
 

src/python/omnimalloc/allocators/random.py

@@ -15,7 +15,7 @@
 class RandomAllocator(GreedyAllocator):
     """Randomized allocator that tries multiple random orders and picks the best."""
 
-    def __init__(self, *, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None:
+    def __init__(self, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None:
         ensure_non_negative(num_trials, "num_trials")
         self._seed = seed
         self._num_trials = num_trials

src/python/omnimalloc/allocators/simulated_annealing.py

@@ -31,7 +31,6 @@ class SimulatedAnnealingAllocator(BaseAllocator):
 
     def __init__(
         self,
-        *,
         seed: int = DEFAULT_SEED,
         max_iterations: int = 3000,
         initial_temperature: float = 3.0,

src/python/omnimalloc/allocators/supermalloc.py

@@ -133,7 +133,6 @@ class SupermallocAllocator(BaseAllocator):
 
     def __init__(
         self,
-        *,
         timeout: float | None = DEFAULT_TIMEOUT,
         heuristics: tuple[Heuristic, ...] = DEFAULT_HEURISTICS,
         num_threads: int | None = None,

src/python/omnimalloc/allocators/tabu_search.py

@@ -30,7 +30,6 @@ class TabuSearchAllocator(BaseAllocator):
 
     def __init__(
         self,
-        *,
         seed: int = DEFAULT_SEED,
         max_iterations: int = 500,
         neighborhood_size: int = 20,

src/python/omnimalloc/allocators/telamalloc.py

@@ -36,7 +36,6 @@ class TelamallocAllocator(BaseAllocator):
 
     def __init__(
         self,
-        *,
         seed: int = DEFAULT_SEED,
         max_backtracks: int = 10000,
         timeout: float | None = DEFAULT_TIMEOUT,

src/python/omnimalloc/analysis/__init__.py

@@ -4,6 +4,10 @@
 
 from ._conflicts import conflict_degrees as conflict_degrees
 from ._conflicts import conflicts as conflicts
+from ._pressure import antichain_pressure as antichain_pressure
+from ._pressure import (
+    antichain_pressure_per_allocation as antichain_pressure_per_allocation,
+)
 from ._pressure import closure_pressure as closure_pressure
 from ._pressure import (
     closure_pressure_per_allocation as closure_pressure_per_allocation,

src/python/omnimalloc/analysis/_pressure.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Sequence
 
-from omnimalloc._cpp import antichain_pressure
+from omnimalloc._cpp import antichain_pressure as _antichain_pressure
 from omnimalloc._cpp import (
     antichain_pressure_per_allocation as _antichain_pressure_per_allocation,
 )
@@ -21,7 +21,7 @@
 from omnimalloc.primitives.utils import ensure_unique_ids
 
 
-def pressure(
+def antichain_pressure(
     allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET
 ) -> int:
     """Peak memory pressure: exact max-weight antichain of the happens-before order.
@@ -36,7 +36,7 @@ def pressure(
     the flow work exceeds `work_budget`; pass `None` to always compute.
     """
     ensure_valid_budget(work_budget)
-    return antichain_pressure(allocations, work_budget)
+    return _antichain_pressure(allocations, work_budget)
 
 
 def closure_pressure(
@@ -46,7 +46,7 @@ def closure_pressure(
 
     C++ enumeration of the join-closure of the birth clocks. Pairwise-
     concurrent allocations need not share a cut, so this can sit strictly
-    below `pressure`; both soundly lower-bound any placement's peak. Raises
+    below `antichain_pressure`; both soundly lower-bound any placement's peak. Raises
     `RuntimeError` once the closure exceeds `closure_cap`; pass `None` to
     always enumerate — memory then grows with the closure itself.
     """
@@ -58,7 +58,7 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int:
     """Peak of a placement: the highest occupied address, max(offset + size).
 
     Simply the pressure the placement realizes after allocation — an upper
-    bound on `pressure` (and so on `closure_pressure`), equal to the max
+    bound on `antichain_pressure` (and so on `closure_pressure`), equal to the max
     entry of `placement_pressure_per_allocation`. Raises `ValueError` on
     unplaced input.
     """
@@ -71,14 +71,14 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int:
     return max(heights, default=0)
 
 
-def pressure_per_allocation(
+def antichain_pressure_per_allocation(
     allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET
 ) -> dict[IdType, int]:
     """Peak pressure over each allocation's own lifetime, keyed by id.
 
     The max-weight antichain through each allocation: the tightest
     order-derived lower bound on the pressure any placement can exhibit
-    while that allocation is live; the max entry equals `pressure`.
+    while that allocation is live; the max entry equals `antichain_pressure`.
     Linearizable instances take one O(N log N) window sweep, genuinely
     partial orders solve one pinned antichain (min flow over the conflict
     neighborhood) per distinct lifetime — exact, but built to certify
@@ -98,7 +98,7 @@ def closure_pressure_per_allocation(
     """Exact realizable peak while each allocation is live, keyed by id.
 
     The max total size at any join-closure cut where the allocation is
-    live. Can sit elementwise strictly below `pressure_per_allocation`,
+    live. Can sit elementwise strictly below `antichain_pressure_per_allocation`,
     since pairwise-concurrent allocations need not share a cut; the max
     entry equals `closure_pressure`. Raises `RuntimeError` once the
     closure exceeds `closure_cap`; pass `None` to always enumerate —
@@ -128,6 +128,11 @@ def placement_pressure_per_allocation(
     return _keyed_by_id(allocations, peaks)
 
 
+# The antichain bound is the canonical/default pressure metric
+pressure = antichain_pressure
+pressure_per_allocation = antichain_pressure_per_allocation
+
+
 def _keyed_by_id(
     allocations: Sequence[Allocation], peaks: list[int]
 ) -> dict[IdType, int]:

src/python/omnimalloc/benchmark/timer.py

@@ -20,7 +20,7 @@ class Timer:
     multiple threads may result in race conditions and inconsistent state.
     """
 
-    def __init__(self, *, auto_start: bool = False) -> None:
+    def __init__(self, auto_start: bool = False) -> None:
         self._start_ns: int | None = None
         self._stop_ns: int | None = None
         self._is_running: bool = False

src/python/omnimalloc/io.py

@@ -3,6 +3,7 @@
 #
 
 import csv
+from collections.abc import Sequence
 from pathlib import Path
 
 from .analysis.clock import time_components
@@ -39,7 +40,7 @@ def _write_pool(pool: Pool, path: Path) -> Path:
     # Any placed allocation brings in the offset column (minimalloc's
     # solution format; unplaced rows leave the cell blank), so save/load
     # round-trips placements instead of silently stripping them.
-    with_offsets = any(alloc.is_allocated for alloc in pool.allocations)
+    with_offsets = pool.any_allocated
     fields = ("id", "lower", "upper", "size")
     if with_offsets:
         fields = (*fields, "offset")
@@ -56,11 +57,12 @@ def _write_pool(pool: Pool, path: Path) -> Path:
 
 
 def save_allocation(
-    entity: System | Memory | Pool, path: str | Path
+    entity: System | Memory | Pool | Sequence[Allocation], path: str | Path
 ) -> tuple[Path, ...]:
     """Save the entity's pools to disk as minimalloc-format CSV files.
 
-    Saving a `Pool` writes exactly `path`; a `Memory` or `System` fans out
+    Saving a `Pool` (or a raw sequence of Allocations, saved as one pool)
+    writes exactly `path`; a `Memory` or `System` fans out
     to one `<stem>_<pool_id>.csv` per pool (the CSV is pool-level minimalloc
     interchange). Returns the tuple of paths actually written. Pools with
     any placed allocation include an `offset` column (minimalloc's solution
@@ -72,10 +74,10 @@ def save_allocation(
     """
     path_ = Path(path)
     path_.parent.mkdir(parents=True, exist_ok=True)
+    if not isinstance(entity, System | Memory | Pool):
+        entity = Pool.from_allocations(entity)
     if isinstance(entity, Pool):
         return (_write_pool(entity, path_),)
-    if not isinstance(entity, (Memory, System)):
-        raise TypeError(f"Unsupported entity type: {type(entity)!r}")
     return tuple(
         _write_pool(pool, path_.with_name(f"{path_.stem}_{name}.csv"))
         for name, pool in _collect_pools(entity).items()

src/python/omnimalloc/primitives/memory.py

@@ -19,19 +19,19 @@
 class Memory:
     """A physical memory unit containing one or more pools.
 
-    `capacity` is the declared memory limit (an input); the computed extent
-    of a placement is `used_size`. `capacity=None` means unbounded.
+    `size` is the declared memory size (an input); the computed extent of a
+    placement is `used_size`. `size=None` means unbounded.
     """
 
     id: IdType
     pools: tuple[Pool, ...]
-    capacity: int | None = None
+    size: int | None = None
 
     def __post_init__(self) -> None:
         if len({pool.id for pool in self.pools}) != len(self.pools):
             raise ValueError("pool ids must be unique")
-        if self.capacity is not None:
-            ensure_non_negative(self.capacity, "capacity")
+        if self.size is not None:
+            ensure_non_negative(self.size, "size")
 
     @cached_property
     def used_size(self) -> int:
@@ -43,9 +43,14 @@ def is_allocated(self) -> bool:
         """True if all pools have been allocated."""
         return all(pool.is_allocated for pool in self.pools)
 
+    @cached_property
+    def any_allocated(self) -> bool:
+        """True if any pool has a placed allocation."""
+        return any(pool.any_allocated for pool in self.pools)
+
     def with_pools(self, pools: tuple[Pool, ...]) -> "Memory":
         """Return new Memory with specified pools."""
-        return Memory(id=self.id, capacity=self.capacity, pools=pools)
+        return Memory(id=self.id, size=self.size, pools=pools)
 
     def allocate(self, allocator: "BaseAllocator") -> "Memory":
         """Apply allocator to all pools."""