From ae45c05e3d92c75d27e2b08997a8cc41d0d4604c Mon Sep 17 00:00:00 2001
From: LTLA <infinite.monkeys.with.keyboards@gmail.com>
Date: Fri, 27 Dec 2024 21:06:53 -0800
Subject: [PATCH] Check that the row indices for each gene set are not out of
 range.

---
 .../aggregate_across_genes.hpp                | 16 +++++++++---
 tests/src/aggregate_across_genes.cpp          | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/include/scran_aggregate/aggregate_across_genes.hpp b/include/scran_aggregate/aggregate_across_genes.hpp
index 0794032..ee58dfc 100644
--- a/include/scran_aggregate/aggregate_across_genes.hpp
+++ b/include/scran_aggregate/aggregate_across_genes.hpp
@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <vector>
 #include <unordered_set>
+#include <stdexcept>
 
 #include "tatami/tatami.hpp"
 
@@ -65,15 +66,22 @@ struct AggregateAcrossGenesResults {
 namespace aggregate_across_genes_internal {
 
 template<typename Index_, typename Gene_, typename Weight_>
-std::vector<Gene_> create_subset(const std::vector<std::tuple<size_t, const Gene_*, const Weight_*> >& gene_sets) {
+std::vector<Gene_> create_subset(const std::vector<std::tuple<size_t, const Gene_*, const Weight_*> >& gene_sets, Index_ nrow) {
     std::unordered_set<Gene_> of_interest;
     for (const auto& set : gene_sets) {
         auto set_size = std::get<0>(set);
         auto set_genes = std::get<1>(set);
         of_interest.insert(set_genes, set_genes + set_size);
     }
+
     std::vector<Index_> subset(of_interest.begin(), of_interest.end());
-    std::sort(subset.begin(), subset.end());
+    if (!subset.empty()) {
+        std::sort(subset.begin(), subset.end());
+        if (subset.front() < 0 || subset.back() >= nrow) {
+            throw std::runtime_error("set indices are out of range");
+        }
+    }
+
     return subset;
 }
 
@@ -97,7 +105,7 @@ void compute_aggregate_by_column(
     const AggregateAcrossGenesOptions& options)
 {
     // Identifying the subset of rows that actually need to be extracted.
-    tatami::VectorPtr<Index_> subset_of_interest = std::make_shared<std::vector<Index_> >(create_subset<Index_>(gene_sets));
+    tatami::VectorPtr<Index_> subset_of_interest = std::make_shared<std::vector<Index_> >(create_subset<Index_>(gene_sets, p.nrow()));
     const auto& subset = *subset_of_interest;
     size_t nsubs = subset.size();
 
@@ -160,7 +168,7 @@ void compute_aggregate_by_row(
     const AggregateAcrossGenesOptions& options)
 {
     // Identifying the subset of rows that actually need to be extracted.
-    auto subset = create_subset<Index_>(gene_sets);
+    auto subset = create_subset<Index_>(gene_sets, p.nrow());
     size_t nsubs = subset.size();
     auto sub_oracle = std::make_shared<tatami::FixedViewOracle<Index_> >(subset.data(), nsubs);
 
diff --git a/tests/src/aggregate_across_genes.cpp b/tests/src/aggregate_across_genes.cpp
index 64f1472..29254db 100644
--- a/tests/src/aggregate_across_genes.cpp
+++ b/tests/src/aggregate_across_genes.cpp
@@ -177,3 +177,28 @@ INSTANTIATE_TEST_SUITE_P(
     AggregateAcrossGenesTest,
     ::testing::Values(1, 3) // number of threads
 );
+
+TEST(AggregateAcrossGenes, OutOfRange) {
+    int nr = 11, nc = 78;
+    auto vec = scran_tests::simulate_vector(nr * nc, []{
+        scran_tests::SimulationParameters sparams;
+        sparams.density = 0.1;
+        return sparams;
+    }());
+
+    tatami::DenseRowMatrix<double, int> mat(nr, nc, std::move(vec));
+    std::vector<int> example { 1, 10, 100 };
+    std::vector<std::tuple<size_t, const int*, const double*> > gene_sets;
+    gene_sets.emplace_back(3, example.data(), static_cast<double*>(NULL));
+
+    scran_aggregate::AggregateAcrossGenesOptions opt;
+    scran_tests::expect_error([&]() {
+        scran_aggregate::aggregate_across_genes(mat, gene_sets, opt);
+    }, "out of range");
+
+    // Also fails if there are negative values.
+    example[0] = -1;
+    scran_tests::expect_error([&]() {
+        scran_aggregate::aggregate_across_genes(mat, gene_sets, opt);
+    }, "out of range");
+}