Explicitly compute number of non-zeros in row

Change-Id: Iaebaf6d23d33dbbbe5a7c7240319bf28cf2bdd3a
diff --git a/internal/ceres/block_sparse_matrix.cc b/internal/ceres/block_sparse_matrix.cc
index 6feedc8..b3d4efd 100644
--- a/internal/ceres/block_sparse_matrix.cc
+++ b/internal/ceres/block_sparse_matrix.cc
@@ -85,23 +85,18 @@
       continue;
     }
 
-    // Compute number of non-zeros per row from number of non-zeros in row-block
-    int row_block_nnz;
+    int row_nnz = 0;
     if constexpr (transpose) {
-      // Transposed block structure comes with nnz filled-in
-      row_block_nnz = row_block.nnz;
+      // Transposed block structure comes with nnz in row-block filled-in
+      row_nnz = row_block.nnz / row_block.block.size;
     } else {
-      // Non-transposed block structure has sequential layout, but nnz field of
-      // block structure is not filled. Difference between values position of
-      // the first and the last cells in row-block gives number of non-zero
-      // elements in row-block excluding the last cell.
-      const auto& front = row_block.cells.front();
-      const auto& back = row_block.cells.back();
-      row_block_nnz =
-          (back.position - front.position) +
-          block_structure->cols[back.block_id].size * row_block.block.size;
+      // Nnz field of non-transposed block structure is not filled and it can
+      // have non-sequential structure (consider the case of jacobian for
+      // Schur-complement solver: E and F blocks are stored separately).
+      for (auto& c : row_block.cells) {
+        row_nnz += cols[c.block_id].size;
+      }
     }
-    const int row_nnz = row_block_nnz / row_block.block.size;
 
     // Row-wise setup of matrix structure
     for (int row = 0; row < row_block.block.size; ++row) {
diff --git a/internal/ceres/block_sparse_matrix_test.cc b/internal/ceres/block_sparse_matrix_test.cc
index f749df3..daa2e2d 100644
--- a/internal/ceres/block_sparse_matrix_test.cc
+++ b/internal/ceres/block_sparse_matrix_test.cc
@@ -113,6 +113,39 @@
       values[i] = i + 1;
     }
     return m;
+  } else if (id == 2) {
+    // Create the following block sparse matrix:
+    // [ 1 2 0 | 6 7 0 ]
+    // [ 3 4 0 | 8 9 0 ]
+    // [ 0 0 5 | 0 0 10]
+    // With cells of the left submatrix preceding cells of the right submatrix
+    CompressedRowBlockStructure* bs = new CompressedRowBlockStructure;
+    bs->cols = {
+        // Block size 2, position 0.
+        Block(2, 0),
+        // Block size 1, position 2.
+        Block(1, 2),
+        // Block size 2, position 3.
+        Block(2, 3),
+        // Block size 1, position 5.
+        Block(1, 5),
+    };
+    bs->rows = {CompressedRow(2), CompressedRow(1)};
+    bs->rows[0].block = Block(2, 0);
+    bs->rows[0].cells = {Cell(0, 0), Cell(2, 5)};
+
+    bs->rows[1].block = Block(1, 2);
+    bs->rows[1].cells = {Cell(1, 4), Cell(3, 9)};
+    auto m = std::make_unique<BlockSparseMatrix>(bs);
+    EXPECT_NE(m, nullptr);
+    EXPECT_EQ(m->num_rows(), 3);
+    EXPECT_EQ(m->num_cols(), 6);
+    EXPECT_EQ(m->num_nonzeros(), 10);
+    double* values = m->mutable_values();
+    for (int i = 0; i < 10; ++i) {
+      values[i] = i + 1;
+    }
+    return m;
   }
   return nullptr;
 }
@@ -477,6 +510,17 @@
     m_expected << 1, 2, 0, 5, 6, 0, 3, 4, 0, 7, 8, 0, 0, 0, 9, 0, 0, 0;
     EXPECT_EQ(m_dense, m_expected);
   }
+
+  {
+    std::unique_ptr<BlockSparseMatrix> m = CreateTestMatrixFromId(2);
+    Matrix m_dense;
+    m->ToDenseMatrix(&m_dense);
+    EXPECT_EQ(m_dense.rows(), 3);
+    EXPECT_EQ(m_dense.cols(), 6);
+    Matrix m_expected(3, 6);
+    m_expected << 1, 2, 0, 6, 7, 0, 3, 4, 0, 8, 9, 0, 0, 0, 5, 0, 0, 10;
+    EXPECT_EQ(m_dense, m_expected);
+  }
 }
 
 TEST(BlockSparseMatrix, ToCRSMatrix) {
@@ -512,6 +556,22 @@
       EXPECT_EQ(m_crs->values()[i], values_expected[i]);
     }
   }
+  {
+    std::unique_ptr<BlockSparseMatrix> m = CreateTestMatrixFromId(2);
+    auto m_crs = m->ToCompressedRowSparseMatrix();
+    std::vector<int> rows_expected = {0, 4, 8, 10};
+    std::vector<int> cols_expected = {0, 1, 3, 4, 0, 1, 3, 4, 2, 5};
+    std::vector<double> values_expected = {1, 2, 6, 7, 3, 4, 8, 9, 5, 10};
+    for (int i = 0; i < rows_expected.size(); ++i) {
+      EXPECT_EQ(m_crs->rows()[i], rows_expected[i]);
+    }
+    for (int i = 0; i < cols_expected.size(); ++i) {
+      EXPECT_EQ(m_crs->cols()[i], cols_expected[i]);
+    }
+    for (int i = 0; i < values_expected.size(); ++i) {
+      EXPECT_EQ(m_crs->values()[i], values_expected[i]);
+    }
+  }
 }
 
 TEST(BlockSparseMatrix, ToCRSMatrixTranspose) {
@@ -551,6 +611,24 @@
       EXPECT_EQ(m_crs_transpose->values()[i], values_expected[i]);
     }
   }
+  {
+    std::unique_ptr<BlockSparseMatrix> m = CreateTestMatrixFromId(2);
+    auto m_crs_transpose = m->ToCompressedRowSparseMatrixTranspose();
+    std::vector<int> rows_expected = {0, 2, 4, 5, 7, 9, 10};
+    std::vector<int> cols_expected = {0, 1, 0, 1, 2, 0, 1, 0, 1, 2};
+    std::vector<double> values_expected = {1, 3, 2, 4, 5, 6, 8, 7, 9, 10};
+    EXPECT_EQ(m_crs_transpose->num_nonzeros(), cols_expected.size());
+    EXPECT_EQ(m_crs_transpose->num_rows(), rows_expected.size() - 1);
+    for (int i = 0; i < rows_expected.size(); ++i) {
+      EXPECT_EQ(m_crs_transpose->rows()[i], rows_expected[i]);
+    }
+    for (int i = 0; i < cols_expected.size(); ++i) {
+      EXPECT_EQ(m_crs_transpose->cols()[i], cols_expected[i]);
+    }
+    for (int i = 0; i < values_expected.size(); ++i) {
+      EXPECT_EQ(m_crs_transpose->values()[i], values_expected[i]);
+    }
+  }
 }
 
 TEST(BlockSparseMatrix, CreateTranspose) {