Update matshell and with_localarray name

Author: Jeremy E Kozdon

src/matshell.jl

@@ -40,13 +40,11 @@ function (::MatOp{PetscLib, PetscInt, LibPETSc.MATOP_MULT})(
     mat = unsafe_pointer_to_objref(ptr)
 
     PetscScalar = PetscLib.PetscScalar
-    x = unsafe_localarray(VecPtr(PetscLib, cx); write = false)
-    y = unsafe_localarray(VecPtr(PetscLib, cy); read = false)
+    x = VecPtr(PetscLib, cx)
+    y = VecPtr(PetscLib, cy)
 
     _mul!(y, mat, x)
 
-    Base.finalize(y)
-    Base.finalize(x)
     return PetscInt(0)
 end
 

src/vec.jl

@@ -406,12 +406,9 @@ function unsafe_localarray(
 end
 
 """
-    with_unsafe_localarray!(
-        f!,
-        x::AbstractVec;
-        read=true,
-        write=true,
-    )
+    with_localarray!(f!, x::AbstractVec; read=true, write=true)
+    with_localarray!(f!, xs...; read=true, write=true)
+    with_localarray!(f!, xs::NTuple{N, AbstractVec}...; read=true, write=true)
 
 Convert `x` to an `Array{PetscScalar}` using [`unsafe_localarray`](@ref) and
 apply the function `f!`.
@@ -427,11 +424,16 @@ end
 !!! note
     `Base.finalize` is automatically called on the array.
 """
-function with_unsafe_localarray!(f!, v::AbstractVec; kwargs...)
-    array = unsafe_localarray(v; kwargs...)
-    f!(array)
-    Base.finalize(array)
+function with_localarray!(f!, vecs::NTuple{N, AbstractVec}; kwargs...) where N
+    arrays = map(vecs) do v
+        unsafe_localarray(v; kwargs...)
+    end
+    f!(arrays...)
+    map(arrays) do array
+        Base.finalize(array)
+    end
 end
+with_localarray!(f!, vecs...; kwargs...) = with_localarray!(f!, vecs; kwargs...)
 
 """
     assemblybegin!(vec::AbstractVec)

test/matshell.jl

@@ -9,7 +9,11 @@ using MPI
 
         local_rows = 10
         local_cols = 5
-        f!(x, y) = x .= [2y; 3y]
+        function f!(p_x, p_y)
+            PETSc.with_localarray!((p_x, p_y)) do x, y
+                x .= [2y; 3y]
+            end
+        end
         x_jl = collect
 
         matshell =

test/mpivec.jl

@@ -59,7 +59,7 @@ using LinearAlgebra: norm
             # 1-based
             @test petsc_x[rng .+ 1] == julia_x
 
-            PETSc.with_unsafe_localarray!(petsc_x) do x
+            PETSc.with_localarray!(petsc_x) do x
                 @test x == julia_x
             end
 

test/vec.jl

@@ -30,7 +30,7 @@ using LinearAlgebra: norm
         @test PETSc.ownershiprange(petsc_x) == 1:N
         @test PETSc.ownershiprange(petsc_x, false) == 0:(N - 1)
 
-        PETSc.with_unsafe_localarray!(petsc_x) do x2
+        PETSc.with_localarray!(petsc_x) do x2
             @test x2 == x
         end
 
@@ -51,7 +51,7 @@ end
         @test PETSc.ownershiprange(petsc_x, false) == 0:(N - 1)
 
         x = rand(PetscScalar, N)
-        PETSc.with_unsafe_localarray!(petsc_x) do x2
+        PETSc.with_localarray!(petsc_x) do x2
             x2 .= x
         end
         @test norm(petsc_x) ≈ norm(x)