From fd1f1ba6ce58c5e33710f2186e969c5ef5d9af32 Mon Sep 17 00:00:00 2001
From: Alexis Montoison <alexis.montoison@polymtl.ca>
Date: Mon, 26 Sep 2022 01:21:40 -0400
Subject: [PATCH] Remove allocations in roots_quadratic

---
 src/krylov_utils.jl | 37 +++++++++++++++++++++----------------
 test/test_aux.jl    | 22 ++++++++++++++++++++++
 2 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/src/krylov_utils.jl b/src/krylov_utils.jl
index b16da57c0..c3dce2817 100644
--- a/src/krylov_utils.jl
+++ b/src/krylov_utils.jl
@@ -111,10 +111,10 @@ function roots_quadratic(q₂ :: T, q₁ :: T, q₀ :: T;
   # Case where q(x) is linear.
   if q₂ == zero(T)
     if q₁ == zero(T)
-      root = [zero(T)]
-      q₀ == zero(T) || (root = T[])
+      root = tuple(zero(T))
+      q₀ == zero(T) || (root = tuple())
     else
-      root = [-q₀ / q₁]
+      root = tuple(-q₀ / q₁)
     end
     return root
   end
@@ -123,26 +123,31 @@ function roots_quadratic(q₂ :: T, q₁ :: T, q₀ :: T;
   rhs = √eps(T) * q₁ * q₁
   if abs(q₀ * q₂) > rhs
     ρ = q₁ * q₁ - 4 * q₂ * q₀
-    ρ < 0 && return T[]
+    ρ < 0 && return tuple()
     d = -(q₁ + copysign(sqrt(ρ), q₁)) / 2
-    roots = [d / q₂, q₀ / d]
+    root1 = d / q₂
+    root2 = q₀ / d
   else
     # Ill-conditioned quadratic.
-    roots = [-q₁ / q₂, zero(T)]
+    root1 = -q₁ / q₂
+    root2 = zero(T)
   end
 
   # Perform a few Newton iterations to improve accuracy.
-  for k = 1 : 2
-    root = roots[k]
-    for it = 1 : nitref
-      q = (q₂ * root + q₁) * root + q₀
-      dq = 2 * q₂ * root + q₁
-      dq == zero(T) && continue
-      root = root - q / dq
-    end
-    roots[k] = root
+  for it = 1 : nitref
+    q = (q₂ * root1 + q₁) * root1 + q₀
+    dq = 2 * q₂ * root1 + q₁
+    dq == zero(T) && continue
+    root1 = root1 - q / dq
+  end
+
+  for it = 1 : nitref
+    q = (q₂ * root2 + q₁) * root2 + q₀
+    dq = 2 * q₂ * root2 + q₁
+    dq == zero(T) && continue
+    root2 = root2 - q / dq
   end
-  return roots
+  return (root1, root2)
 end
 
 
diff --git a/test/test_aux.jl b/test/test_aux.jl
index 5ac2b401c..72815ff2f 100644
--- a/test/test_aux.jl
+++ b/test/test_aux.jl
@@ -36,54 +36,76 @@
   @testset "roots_quadratic" begin
     # test roots of a quadratic
     roots = Krylov.roots_quadratic(0.0, 0.0, 0.0)
+    allocations = @allocated Krylov.roots_quadratic(0.0, 0.0, 0.0)
     @test length(roots) == 1
     @test roots[1] == 0.0
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(0.0, 0.0, 1.0)
+    allocations = @allocated Krylov.roots_quadratic(0.0, 0.0, 1.0)
     @test length(roots) == 0
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(0.0, 3.14, -1.0)
+    allocations = @allocated Krylov.roots_quadratic(0.0, 3.14, -1.0)
     @test length(roots) == 1
     @test roots[1] == 1.0 / 3.14
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(1.0, 0.0, 1.0)
+    allocations = @allocated Krylov.roots_quadratic(1.0, 0.0, 1.0)
     @test length(roots) == 0
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(1.0, 0.0, 0.0)
+    allocations = @allocated Krylov.roots_quadratic(1.0, 0.0, 0.0)
     @test length(roots) == 2
     @test roots[1] == 0.0
     @test roots[2] == 0.0
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(1.0, 3.0, 2.0)
+    allocations = @allocated Krylov.roots_quadratic(1.0, 3.0, 2.0)
     @test length(roots) == 2
     @test roots[1] ≈ -2.0
     @test roots[2] ≈ -1.0
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(1.0e+8, 1.0, 1.0)
+    allocations = @allocated Krylov.roots_quadratic(1.0e+8, 1.0, 1.0)
     @test length(roots) == 0
+    @test allocations == 0
 
     # ill-conditioned quadratic
     roots = Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=0)
+    allocations = @allocated Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=0)
     @test length(roots) == 2
     @test roots[1] == 1.0e+13
     @test roots[2] == 0.0
+    @test allocations == 0
 
     # iterative refinement is crucial!
     roots = Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=1)
+    allocations = @allocated Krylov.roots_quadratic(-1.0e-8, 1.0e+5, 1.0, nitref=1)
     @test length(roots) == 2
     @test roots[1] == 1.0e+13
     @test roots[2] == -1.0e-05
+    @test allocations == 0
 
     # not ill-conditioned quadratic
     roots = Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=0)
+    allocations = @allocated Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=0)
     @test length(roots) == 2
     @test isapprox(roots[1],  1.0e+7, rtol=1.0e-6)
     @test isapprox(roots[2], -1.0, rtol=1.0e-6)
+    @test allocations == 0
 
     roots = Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=1)
+    allocations = @allocated Krylov.roots_quadratic(-1.0e-7, 1.0, 1.0, nitref=1)
     @test length(roots) == 2
     @test isapprox(roots[1], 1.0e+7, rtol=1.0e-6)
     @test isapprox(roots[2], -1.0, rtol=1.0e-6)
+    @test allocations == 0
   end
 
   @testset "to_boundary" begin