From 04ced10feef2cdb1a90a9690a0f3916c0f84ba8f Mon Sep 17 00:00:00 2001
From: Sacha Verweij <sacha@stanford.edu>
Date: Sat, 4 Mar 2017 12:39:57 -0800
Subject: [PATCH] Fixes and stronger tests for Tridiagonal setindex!

Makes Tridiagonal setindex! support assigning zero off the main, sub, and super diagonals.
Strengthens tests for Tridiagonal setindex!.
---
 base/linalg/tridiag.jl | 13 ++++++++-----
 test/linalg/tridiag.jl | 13 +++++++------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/base/linalg/tridiag.jl b/base/linalg/tridiag.jl
index 279e2aca796ad..805f5516f1354 100644
--- a/base/linalg/tridiag.jl
+++ b/base/linalg/tridiag.jl
@@ -559,15 +559,18 @@ function getindex{T}(A::Tridiagonal{T}, i::Integer, j::Integer)
 end
 
 function setindex!(A::Tridiagonal, x, i::Integer, j::Integer)
+    @boundscheck checkbounds(A, i, j)
     if i == j
-        A.d[i] = x
+        @inbounds A.d[i] = x
     elseif i - j == 1
-        A.dl[j] = x
+        @inbounds A.dl[j] = x
     elseif j - i == 1
-        A.du[i] = x
-    else
-        throw(ArgumentError("cannot set elements outside the sub, main, or super diagonals"))
+        @inbounds A.du[i] = x
+    elseif !iszero(x)
+        throw(ArgumentError(string("cannot set entry ($i, $j) off ",
+            "the tridiagonal band to a nonzero value ($x)")))
     end
+    return x
 end
 
 ## structured matrix methods ##
diff --git a/test/linalg/tridiag.jl b/test/linalg/tridiag.jl
index d6f12d402915d..98371fff8a3b0 100644
--- a/test/linalg/tridiag.jl
+++ b/test/linalg/tridiag.jl
@@ -459,12 +459,13 @@ let n = 12 #Size of matrix problem to test
         @test_throws BoundsError A[1,n+1]
 
         debug && println("setindex!")
-        @test_throws ArgumentError A[n,1] = 1
-        @test_throws ArgumentError A[1,n] = 1
-        A[3,3] = A[3,3]
-        A[2,3] = A[2,3]
-        A[3,2] = A[3,2]
-        @test A == fA
+        @test_throws BoundsError A[n + 1, 1] = 0 # test bounds check
+        @test_throws BoundsError A[1, n + 1] = 0 # test bounds check
+        @test (A[3, 3] = A[3, 3]; A == fA) # test assignment on the main diagonal
+        @test (A[3, 2] = A[3, 2]; A == fA) # test assignment on the subdiagonal
+        @test (A[2, 3] = A[2, 3]; A == fA) # test assignment on the superdiagonal
+        @test ((A[1, 3] = 0) == 0; A == fA) # test zero assignment off the main/sub/super diagonal
+        @test_throws ArgumentError A[1, 3] = 1 # test non-zero assignment off the main/sub/super diagonal
     end
 end