From c6f398122db15eae79be7fdf00958b4b370d0d43 Mon Sep 17 00:00:00 2001
From: Piotr Rzysko <piotr.rzysko@gmail.com>
Date: Wed, 30 Jul 2025 20:28:24 +0200
Subject: [PATCH] Add validation for constraint usage length returned by
 best_index

Additional changes:
- Update IndexInfo documentation to clarify that constraint_usages must
  have exact 1:1 correspondence with input ConstraintInfo array. The code
  translating constraints into VFilter arguments heavily relies on this.
- Fix best_index implementation in test extension to comply with new
  validation requirements by returning usage entry for each constraint
---
 core/translate/main_loop.rs  | 10 +++++-
 extensions/core/src/vtabs.rs |  2 ++
 extensions/tests/src/lib.rs  | 64 +++++++++++++++++++++---------------
 3 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs
index 7c414b0d8..0663b813a 100644
--- a/core/translate/main_loop.rs
+++ b/core/translate/main_loop.rs
@@ -14,7 +14,7 @@ use crate::{
         insn::{CmpInsFlags, IdxInsertFlags, Insn},
         BranchOffset, CursorID,
     },
-    Result,
+    LimboError, Result,
 };
 
 use super::{
@@ -483,6 +483,14 @@ pub fn open_loop(
                             // maybe encode more info on t_ctx? we need: [col_idx, is_descending]
                             let index_info = vtab.best_index(&converted_constraints, &[]);
 
+                            if index_info.constraint_usages.len() != converted_constraints.len() {
+                                return Err(LimboError::ExtensionError(format!(
+                                    "Constraint usage count mismatch (expected {}, got {})",
+                                    converted_constraints.len(),
+                                    index_info.constraint_usages.len()
+                                )));
+                            }
+
                             // Determine the number of VFilter arguments (constraints with an argv_index).
                             let args_needed = index_info
                                 .constraint_usages
diff --git a/extensions/core/src/vtabs.rs b/extensions/core/src/vtabs.rs
index 95af9a9ce..2b041a8af 100644
--- a/extensions/core/src/vtabs.rs
+++ b/extensions/core/src/vtabs.rs
@@ -222,6 +222,8 @@ pub struct IndexInfo {
     /// Estimated number of rows that the query will return
     pub estimated_rows: u32,
     /// List of constraints that can be used to optimize the query.
+    /// Each `ConstraintInfo` passed to `best_index` must have a corresponding entry in `constraint_usages`.
+    /// The length and order are important—they must exactly match the input `ConstraintInfo` array.
     pub constraint_usages: Vec<ConstraintUsage>,
 }
 impl Default for IndexInfo {
diff --git a/extensions/tests/src/lib.rs b/extensions/tests/src/lib.rs
index 88ff9f26a..a7b0939b0 100644
--- a/extensions/tests/src/lib.rs
+++ b/extensions/tests/src/lib.rs
@@ -168,43 +168,55 @@ impl VTable for KVStoreTable {
     }
 
     fn best_index(constraints: &[ConstraintInfo], _order_by: &[OrderByInfo]) -> IndexInfo {
+        let mut constraint_usages = Vec::with_capacity(constraints.len());
+        let mut idx_num = -1;
+        let mut idx_str = None;
+        let mut estimated_cost = 1000.0;
+        let mut estimated_rows = u32::MAX;
+
         // Look for: key = ?
         for constraint in constraints.iter() {
             if constraint.usable
                 && constraint.op == ConstraintOp::Eq
                 && constraint.column_index == 1
+                && idx_num == -1
+            // Only use the first usable key = ? constraint
             {
-                // this extension wouldn't support order by but for testing purposes,
-                // we will consume it if we find an ASC order by clause on the value column
-                let mut consumed = false;
-                if let Some(order) = _order_by.first() {
-                    if order.column_index == 2 && !order.desc {
-                        consumed = true;
-                    }
-                }
+                constraint_usages.push(ConstraintUsage {
+                    omit: true,
+                    argv_index: Some(1),
+                });
+                idx_num = 1;
+                idx_str = Some("key_eq".to_string());
+                estimated_cost = 10.0;
+                estimated_rows = 4;
                 log::debug!("xBestIndex: constraint found for 'key = ?'");
-                return IndexInfo {
-                    idx_num: 1,
-                    idx_str: Some("key_eq".to_string()),
-                    order_by_consumed: consumed,
-                    estimated_cost: 10.0,
-                    estimated_rows: 4,
-                    constraint_usages: vec![ConstraintUsage {
-                        omit: true,
-                        argv_index: Some(1),
-                    }],
-                };
+            } else {
+                constraint_usages.push(ConstraintUsage {
+                    omit: false,
+                    argv_index: None,
+                });
             }
         }
 
-        // fallback: full scan
-        log::debug!("No usable constraints found, using full scan");
+        // this extension wouldn't support order by but for testing purposes,
+        // we will consume it if we find an ASC order by clause on the value column
+        let order_by_consumed = idx_num == 1
+            && _order_by
+                .first()
+                .is_some_and(|order| order.column_index == 2 && !order.desc);
+
+        if idx_num == -1 {
+            log::debug!("No usable constraints found, using full scan");
+        }
+
         IndexInfo {
-            idx_num: -1,
-            idx_str: None,
-            order_by_consumed: false,
-            estimated_cost: 1000.0,
-            ..Default::default()
+            idx_num,
+            idx_str,
+            order_by_consumed,
+            estimated_cost,
+            estimated_rows,
+            constraint_usages,
         }
     }