refactor: Split ArrayBorrowCompiler off from ArrayGetitemCompiler

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,14 +313,24 @@ 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)
 
     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."""
 

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)