feat: better sub for [ZZ|QQ]Matrix

Author: fieker

src/matrix.jl

@@ -20,6 +20,30 @@ function sub(x::_MatTypes, r::AbstractUnitRange{Int}, c::AbstractUnitRange{Int})
   return deepcopy(view(x, r, c))
 end
 
+# flint does not support windows with non-continuous row or column indices
+# we overload some specific types to not hit the AA fallback
+for T in [ZZMatrix, QQMatrix]
+  @eval begin
+    # have to "duplicate" this, otherwise we get ambiguity errors
+    function sub(x::$T, r::AbstractUnitRange{Int}, c::AbstractUnitRange{Int})
+      return deepcopy(view(x, r, c))
+    end
+
+    function sub(M::$T, r::AbstractVector{Int}, c::AbstractVector{Int})
+      N = zero_matrix(base_ring(M), length(r), length(c))
+      GC.@preserve M N begin
+        for (ii, i) in enumerate(r)
+          for (jj, j) in enumerate(c)
+            set!(mat_entry_ptr(N, ii, jj), mat_entry_ptr(M, i, j))
+          end
+        end
+      end
+      return N
+    end
+  end
+end
+
+# make x[r, c] sugar for sub(x, r, c)
 getindex(x::_MatTypes, r::AbstractUnitRange{Int}, c::AbstractUnitRange{Int}) = sub(x, r, c)
 
 ################################################################################

test/flint/fmpz_mat-test.jl

@@ -259,6 +259,12 @@ end
   C[1, 1] = 20
   @test B == matrix_space(ZZ, 2, 2)([10 2; 4 5])
   @test A == S([1 2 3; 4 5 6; 7 8 9])
+
+  m = ZZ[1 2 4; 5 6 7]
+  n = m[[1], [1, 3]]
+  @test n == ZZ[1 4]
+  n = m[1:2, 1:2]
+  @test n == ZZ[1 2; 5 6]
 end
 
 @testset "ZZMatrix.unary_ops" begin

test/flint/fmpz_mod_mat-test.jl

@@ -858,6 +858,11 @@ end
   @test collect(v) == [3, 6, 9]
   v[1] = 1
   @test A == S([1 2 1; 4 5 6; 7 8 9])
+
+  m = Z17[1 2 4; 5 6 7]
+  n = m[[1], [1, 3]]
+  @test n == Z17[1 4]
+  n = m[1:2, 1:2]
 end
 
 @testset "ZZModMatrix.sub" begin