fixed ASAN check on rmm-tree

Author: RSMT98

.github/workflows/build-test.yml

@@ -1,46 +1,45 @@
-name: Tests (Asan)
-
-on:
-  push:
-    branches: [ main ]
-  pull_request:
-    branches: [ main ]
-
-jobs:
-  build-and-test:
-    runs-on: ubuntu-latest
-    
-    steps:
-    - uses: actions/checkout@v4
-
-    - name: Create Build Directory
-      run: mkdir build
-
-    - name: Configure CMake
-      working-directory: ./build
-      run: cmake -DDISABLE_AVX512=ON -DENABLE_ADDRESS_SANITIZER=ON -DPIXIE_BENCHMARKS=OFF ..
-
-    - name: Build Project
-      working-directory: ./build
-      run: make -j
-
-    - name: Run Unittests
-      working-directory: ./build
-      run: ./unittests
-    
-    - name: Run LOUDS Tree Tests
-      working-directory: ./build
-      run: ./louds_tree_tests
-
-    - name: Run Benchmark Tests
-      working-directory: ./build
-      run: ./benchmark_tests
-
-    #   TODO: fix RmM tests under Asan
-    # - name: Run RmM Tree Tests
-    #   working-directory: ./build
-    #   run: ./test_rmm
-  
+name: Tests (Asan)
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build-and-test:
+    runs-on: ubuntu-latest
+    
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Create Build Directory
+      run: mkdir build
+
+    - name: Configure CMake
+      working-directory: ./build
+      run: cmake -DDISABLE_AVX512=ON -DENABLE_ADDRESS_SANITIZER=ON -DPIXIE_BENCHMARKS=OFF ..
+
+    - name: Build Project
+      working-directory: ./build
+      run: make -j
+
+    - name: Run Unittests
+      working-directory: ./build
+      run: ./unittests
+    
+    - name: Run LOUDS Tree Tests
+      working-directory: ./build
+      run: ./louds_tree_tests
+
+    - name: Run Benchmark Tests
+      working-directory: ./build
+      run: ./benchmark_tests
+
+    - name: Run RmM Tree Tests
+      working-directory: ./build
+      run: ./test_rmm
+  
   build-and-test-with-SDE:
     runs-on: ubuntu-latest
     timeout-minutes: 60
@@ -92,4 +91,4 @@ jobs:
           exit 0
         fi
         exit $rc
-
+

include/pixie/rmm_tree.h

@@ -270,26 +270,41 @@ class RmMTree {
     if (node_pattern10_count[node_index] < target_pattern_rank) {
       return npos;
     }
+    const size_t tree_size = segment_size_bits.size() - 1;
     size_t segment_base = 0;
     while (node_index < first_leaf_index) {
-      const size_t left_child = node_index << 1, right_child = left_child | 1;
+      const size_t left_child = node_index << 1;
+      const size_t left_segment_size =
+          (left_child <= tree_size) ? segment_size_bits[left_child] : 0;
+      if (left_segment_size == 0) {
+        return npos;
+      }
+
+      const size_t left_count = node_pattern10_count[left_child];
+      if (left_count >= target_pattern_rank) {
+        node_index = left_child;
+        continue;
+      }
+
+      size_t remaining_rank = target_pattern_rank - left_count;
+      const size_t right_child = left_child | 1;
+      const bool has_right =
+          (right_child <= tree_size) && (segment_size_bits[right_child] != 0);
+      if (!has_right) {
+        return npos;
+      }
+
       const size_t crossing_pattern =
           (node_last_bit[left_child] == 1 && node_first_bit[right_child] == 0)
               ? 1u
               : 0u;
-      if (node_pattern10_count[left_child] >= target_pattern_rank) {
-        node_index = left_child;
-        continue;
-      }
-      size_t remaining_rank =
-          target_pattern_rank - node_pattern10_count[left_child];
       if (crossing_pattern) {
         if (remaining_rank == 1) {
-          return segment_base + segment_size_bits[left_child] - 1;
+          return segment_base + left_segment_size - 1;
         }
         --remaining_rank;
       }
-      segment_base += segment_size_bits[left_child];
+      segment_base += left_segment_size;
       node_index = right_child;
       target_pattern_rank = remaining_rank;
     }

src/tests/test_rmm.cpp

@@ -602,6 +602,27 @@ TEST(RmMEdgeCases, PartialLastLeafSelects) {
   expect_rank_select_equal(rm_select0, nv_select0, n);
 }
 
+TEST(RmMEdgeCases, Select10OnIncompleteInternalNode) {
+  constexpr size_t leaf_block_bits = 256;
+  const size_t n = (leaf_block_bits * 2) + 32;  // exactly 3 leaves
+  std::string bits(n, '1');
+
+  // Put all "10" patterns into the last (partial) leaf.
+  for (size_t i = leaf_block_bits * 2; i + 1 < n; i += 4) {
+    bits[i] = '1';
+    bits[i + 1] = '0';
+  }
+
+  pixie::RmMTree rm(bits, leaf_block_bits);
+  NaiveRmM nv(bits);
+
+  const size_t pairs10 = nv.rank10(n);
+  ASSERT_GT(pairs10, 0u);
+  for (size_t k = 1; k <= pairs10 + 1; ++k) {
+    EXPECT_EQ(rm.select10(k), nv.select10(k)) << "select10 k=" << k;
+  }
+}
+
 /**
  * Invalid arguments should fail fast and return npos/0 as specified.
  * Covers bad ranks, bad ranges and out-of-bounds BP navigation calls.