From b7883a8ef18b2ab1b6bd2ef7cd238282f88e95b3 Mon Sep 17 00:00:00 2001
From: Tim Quatmann <tim.quatmann@cs.rwth-aachen.de>
Date: Wed, 12 Aug 2020 15:59:50 +0200
Subject: [PATCH] Used the new SparseMatrixBuilder::addDiagonalEntry to
 simplify some of the LRA code.

---
 ...arseDeterministicInfiniteHorizonHelper.cpp | 20 +++-------------
 .../infinitehorizon/internal/LraViHelper.cpp  | 23 +++----------------
 2 files changed, 6 insertions(+), 37 deletions(-)

diff --git a/src/storm/modelchecker/helper/infinitehorizon/SparseDeterministicInfiniteHorizonHelper.cpp b/src/storm/modelchecker/helper/infinitehorizon/SparseDeterministicInfiniteHorizonHelper.cpp
index f7327d8c0..511a67787 100644
--- a/src/storm/modelchecker/helper/infinitehorizon/SparseDeterministicInfiniteHorizonHelper.cpp
+++ b/src/storm/modelchecker/helper/infinitehorizon/SparseDeterministicInfiniteHorizonHelper.cpp
@@ -169,15 +169,13 @@ namespace storm {
                         builder.addNextValue(row, 0, -storm::utility::one<ValueType>());
                     }
                     // Compute weighted sum over successor state. As this is a BSCC, each successor state will again be in the BSCC.
-                    auto diagonalValue = storm::utility::zero<ValueType>();
                     if (row > 0) {
                         if (isEquationSystemFormat) {
-                            diagonalValue = rateAtState;
-                        } else {
-                            diagonalValue = storm::utility::one<ValueType>() - rateAtState;
+                            builder.addDiagonalEntry(row, rateAtState);
+                        } else if (!storm::utility::isOne(rateAtState)) {
+                            builder.addDiagonalEntry(row, storm::utility::one<ValueType>() - rateAtState);
                         }
                     }
-                    bool needDiagonalEntry = !storm::utility::isZero(diagonalValue);
                     for (auto const& entry : this->_transitionMatrix.getRow(globalState)) {
                         uint64_t col = toLocalIndexMap[entry.getColumn()];
                         if (col == 0) {
@@ -188,20 +186,8 @@ namespace storm {
                         if (isEquationSystemFormat) {
                             entryValue = -entryValue;
                         }
-                        if (needDiagonalEntry && col >= row) {
-                            if (col == row) {
-                                entryValue += diagonalValue;
-                            } else { // col > row
-                                builder.addNextValue(row, row, diagonalValue);
-                            }
-                            needDiagonalEntry = false;
-                        }
                         builder.addNextValue(row, col, entryValue);
                     }
-                    if (needDiagonalEntry) {
-                        builder.addNextValue(row, row, diagonalValue);
-                    }
-
                     eqSysVector.push_back(stateValuesGetter(globalState) + rateAtState * actionValuesGetter(globalState));
                     ++row;
                 }
diff --git a/src/storm/modelchecker/helper/infinitehorizon/internal/LraViHelper.cpp b/src/storm/modelchecker/helper/infinitehorizon/internal/LraViHelper.cpp
index c92618735..e5ec57c6f 100644
--- a/src/storm/modelchecker/helper/infinitehorizon/internal/LraViHelper.cpp
+++ b/src/storm/modelchecker/helper/infinitehorizon/internal/LraViHelper.cpp
@@ -87,43 +87,26 @@ namespace storm {
                                     tsToIsTransitionsBuilder.newRowGroup(currTsRow);
                                 }
                             }
-                            // We need to uniformize which means that a diagonal entry for the selfloop will be inserted.
                             // If there are exit rates, the uniformization factor needs to be updated.
                             if (exitRates) {
                                 uniformizationFactor = (*exitRates)[componentState] / _uniformizationRate;
                             }
+                            // We need to uniformize which means that a diagonal entry for the selfloop will be inserted.
                             ValueType selfLoopProb = storm::utility::one<ValueType>() - uniformizationFactor;
-                            uint64_t selfLoopColumn = toSubModelStateMapping[componentState];
                             for (auto componentChoiceIt = getComponentElementChoicesBegin(element); componentChoiceIt != getComponentElementChoicesEnd(element); ++componentChoiceIt) {
-                                bool insertedDiagElement = false;
+                                tsTransitionsBuilder.addDiagonalEntry(currTsRow, selfLoopProb);
                                 for (auto const& entry : this->_transitionMatrix.getRow(*componentChoiceIt)) {
                                     uint64_t subModelColumn = toSubModelStateMapping[entry.getColumn()];
                                     if (isTimedState(entry.getColumn())) {
                                         // We have a transition from a timed state to a timed state
                                         STORM_LOG_ASSERT(subModelColumn < numTsSubModelStates, "Invalid state for timed submodel");
-                                        if (!insertedDiagElement && subModelColumn > selfLoopColumn) {
-                                            // We passed the diagonal entry, so add it now before moving on to the next entry
-                                            tsTransitionsBuilder.addNextValue(currTsRow, selfLoopColumn, selfLoopProb);
-                                            insertedDiagElement = true;
-                                        }
-                                        if (!insertedDiagElement && subModelColumn == selfLoopColumn) {
-                                            // The current entry is the diagonal (selfloop) entry
-                                            tsTransitionsBuilder.addNextValue(currTsRow, selfLoopColumn, selfLoopProb + uniformizationFactor * entry.getValue());
-                                            insertedDiagElement = true;
-                                        } else {
-                                            // The diagonal element either has been inserted already or still lies in front
-                                            tsTransitionsBuilder.addNextValue(currTsRow, subModelColumn,  uniformizationFactor * entry.getValue());
-                                        }
+                                        tsTransitionsBuilder.addNextValue(currTsRow, subModelColumn,  uniformizationFactor * entry.getValue());
                                     } else {
                                         // We have a transition from a timed to a instant state
                                         STORM_LOG_ASSERT(subModelColumn < numIsSubModelStates, "Invalid state for instant submodel");
                                         tsToIsTransitionsBuilder.addNextValue(currTsRow, subModelColumn, uniformizationFactor * entry.getValue());
                                     }
                                 }
-                                // If the diagonal entry for the MS matrix still has not been set, we do that now
-                                if (!insertedDiagElement) {
-                                    tsTransitionsBuilder.addNextValue(currTsRow, selfLoopColumn, selfLoopProb);
-                                }
                                 ++currTsRow;
                             }
                         } else {