From c1bc72805e89f5d1097b33c54edda19c1203a49a Mon Sep 17 00:00:00 2001 From: Alan Lawrence Date: Thu, 16 Oct 2025 16:24:58 +0100 Subject: [PATCH] Split ArrayBorrowCompiler off from ArrayGetitemCompiler --- .../std/_internal/compiler/array.py | 20 +++++++++++++++---- guppylang/src/guppylang/std/array.py | 7 ++++--- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/guppylang-internals/src/guppylang_internals/std/_internal/compiler/array.py b/guppylang-internals/src/guppylang_internals/std/_internal/compiler/array.py index 2bcde6ec6..029e58a27 100644 --- a/guppylang-internals/src/guppylang_internals/std/_internal/compiler/array.py +++ b/guppylang-internals/src/guppylang_internals/std/_internal/compiler/array.py @@ -276,7 +276,9 @@ def compile(self, args: list[Wire]) -> list[Wire]: class ArrayGetitemCompiler(ArrayCompiler): - """Compiler for the `array.__getitem__` function.""" + """Compiler for the `array.__getitem__` function, used for both classical and + linear arrays e.g. when pulling out arguments to pass to a function. + (In the linear, non-owned, case they will be replaced afterwards.)""" def _build_classical_getitem(self, array: Wire, idx: Wire) -> CallReturnWires: """Constructs `__getitem__` for classical arrays.""" @@ -293,8 +295,8 @@ def _build_classical_getitem(self, array: Wire, idx: Wire) -> CallReturnWires: inout_returns=[arr], ) - def _build_linear_getitem(self, array: Wire, idx: Wire) -> CallReturnWires: - """Constructs `array.__getitem__` for linear arrays.""" + def _build_borrow(self, array: Wire, idx: Wire) -> CallReturnWires: + """Constructs `borrow_array.borrow` for linear-typed arrays.""" idx = self.builder.add_op(convert_itousize(), idx) arr, elem = self.builder.add_op( barray_borrow(self.elem_ty, self.length), @@ -311,7 +313,7 @@ def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires: [elem_ty_arg, _] = self.type_args assert isinstance(elem_ty_arg, TypeArg) if not elem_ty_arg.ty.copyable: - return self._build_linear_getitem(array, idx) + return self._build_borrow(array, idx) else: return self._build_classical_getitem(array, idx) @@ -319,6 +321,16 @@ def compile(self, args: list[Wire]) -> list[Wire]: raise InternalGuppyError("Call compile_with_inouts instead") +class ArrayBorrowCompiler(ArrayGetitemCompiler): + """Compiler for the `array.__borrow__` function only, used for ArrayIter.""" + + def _build_classical_getitem(self, array: Wire, idx: Wire) -> CallReturnWires: + raise InternalGuppyError("ArrayIter is Hugr-polymorphic with Linear TypeParam") + # Note that even if we monomorphized in guppy and needed to handle copyable + # here, we could still use just a borrow (not the more complex + # _build_classical_getitem), as ArrayIter ends with a discard-all-borrowed. + + class ArraySetitemCompiler(ArrayCompiler): """Compiler for the `array.__setitem__` function.""" diff --git a/guppylang/src/guppylang/std/array.py b/guppylang/src/guppylang/std/array.py index b4fde115a..9a823cee0 100644 --- a/guppylang/src/guppylang/std/array.py +++ b/guppylang/src/guppylang/std/array.py @@ -16,6 +16,7 @@ from guppylang_internals.definition.custom import CopyInoutCompiler from guppylang_internals.std._internal.checker import ArrayCopyChecker, NewArrayChecker from guppylang_internals.std._internal.compiler.array import ( + ArrayBorrowCompiler, ArrayDiscardAllUsedCompiler, ArrayGetitemCompiler, ArraySetitemCompiler, @@ -101,7 +102,7 @@ def __next__( self: ArrayIter[L, n] @ owned, ) -> Option[tuple[L, ArrayIter[L, n]]]: if self.i < int(n): - elem = _array_unsafe_getitem(self.xs, self.i) + elem = _barray_unsafe_borrow(self.xs, self.i) return some((elem, ArrayIter(self.xs, self.i + 1))) _array_discard_all_used(self.xs) return nothing() @@ -111,8 +112,8 @@ def __next__( def _array_discard_all_used(xs: array[L, n] @ owned) -> None: ... -@custom_function(ArrayGetitemCompiler()) -def _array_unsafe_getitem(xs: array[L, n], idx: int) -> L: ... +@custom_function(ArrayBorrowCompiler()) +def _barray_unsafe_borrow(xs: array[L, n], idx: int) -> L: ... @extend_type(frozenarray_type_def)