Repository: tursodatabase/limbo (rust)

Issue:
bug: limbo tries to use an index on an expression like t.y > t.x (comparing two cols of same table)
```sql
limbo> .schema
CREATE TABLE t(x primary key, y);

limbo> explain select * from t where y > x;
addr  opcode             p1    p2    p3    p4             p5  comment
----  -----------------  ----  ----  ----  -------------  --  -------
0     Init               0     16    0                    0   Start at 16
1     OpenReadAsync      0     2     0                    0   table=t, root=2
2     OpenReadAwait      0     0     0                    0   
3     OpenReadAsync      1     3     0                    0   table=sqlite_autoindex_t_1, root=3
4     OpenReadAwait      0     0     0                    0   
5     Null               0     3     0                    0   r[3]=NULL
6     SeekGT             1     15    3                    0   
7       Column           0     1     4                    0   r[4]=t.y
8       IdxGE            1     15    4                    0   
9       DeferredSeek     1     0     0                    0   
10      Column           0     0     1                    0   r[1]=t.x
11      Column           0     1     2                    0   r[2]=t.y
12      ResultRow        1     2     0                    0   output=r[1..2]
13    NextAsync          1     0     0                    0   
14    NextAwait          1     7     0                    0   
15    Halt               0     0     0                    0   
16    Transaction        0     0     0                    0   write=false
17    Goto               0     1     0                    0 
```

This is completely incorrect lol. I will probably shove a fix for this into #1288 or make a new PR based on that branch, because that one refactors index selection logic.

Pull request:
Fix two issues with indexes
Fixes #1298

- Fixes Limbo trying to use an index using a WHERE clause constraint that refers to the same table on both sides, e.g. `WHERE t.x = t.x`
- Fixes not using indexes when the relevant expression is paren wrapped, e.g.
    - `SELECT * FROM t WHERE (indexcol) > 5`
    - `SELECT * FROM t WHERE (indexcol > 5)`
- Changes existing table logical expr fuzz test to have primary keys (which helped me find both issues above)

Patch:
diff --git a/core/translate/optimizer.rs b/core/translate/optimizer.rs
index 49543fd6..7c29b883 100644
--- a/core/translate/optimizer.rs
+++ b/core/translate/optimizer.rs
@@ -10,9 +10,12 @@ use crate::{
     Result,
 };
 
-use super::plan::{
-    DeletePlan, Direction, GroupBy, IterationDirection, Operation, Plan, Search, SeekDef, SeekKey,
-    SelectPlan, TableReference, UpdatePlan, WhereTerm,
+use super::{
+    plan::{
+        DeletePlan, Direction, EvalAt, GroupBy, IterationDirection, Operation, Plan, Search,
+        SeekDef, SeekKey, SelectPlan, TableReference, UpdatePlan, WhereTerm,
+    },
+    planner::determine_where_to_eval_expr,
 };
 
 pub fn optimize_plan(plan: &mut Plan, schema: &Schema) -> Result<()> {
@@ -794,6 +797,47 @@ enum BinaryExprSide {
     Rhs,
 }
 
+/// Recursively unwrap parentheses from an expression
+/// e.g. (((t.x > 5))) -> t.x > 5
+fn unwrap_parens<T>(expr: T) -> Result<T>
+where
+    T: UnwrapParens,
+{
+    expr.unwrap_parens()
+}
+
+trait UnwrapParens {
+    fn unwrap_parens(self) -> Result<Self>
+    where
+        Self: Sized;
+}
+
+impl UnwrapParens for &ast::Expr {
+    fn unwrap_parens(self) -> Result<Self> {
+        match self {
+            ast::Expr::Column { .. } => Ok(self),
+            ast::Expr::Parenthesized(exprs) => match exprs.len() {
+                1 => unwrap_parens(exprs.first().unwrap()),
+                _ => crate::bail_parse_error!("expected single expression in parentheses"),
+            },
+            _ => Ok(self),
+        }
+    }
+}
+
+impl UnwrapParens for ast::Expr {
+    fn unwrap_parens(self) -> Result<Self> {
+        match self {
+            ast::Expr::Column { .. } => Ok(self),
+            ast::Expr::Parenthesized(mut exprs) => match exprs.len() {
+                1 => unwrap_parens(exprs.pop().unwrap()),
+                _ => crate::bail_parse_error!("expected single expression in parentheses"),
+            },
+            _ => Ok(self),
+        }
+    }
+}
+
 /// Get the position of a column in an index
 /// For example, if there is an index on table T(x,y) then y's position in the index is 1.
 fn get_column_position_in_index(
@@ -801,20 +845,20 @@ fn get_column_position_in_index(
     table_index: usize,
     table_reference: &TableReference,
     index: &Arc<Index>,
-) -> Option<usize> {
-    let ast::Expr::Column { table, column, .. } = expr else {
-        return None;
+) -> Result<Option<usize>> {
+    let ast::Expr::Column { table, column, .. } = unwrap_parens(expr)? else {
+        return Ok(None);
     };
     if *table != table_index {
-        return None;
+        return Ok(None);
     }
     let Some(column) = table_reference.table.get_column_at(*column) else {
-        return None;
+        return Ok(None);
     };
-    index
+    Ok(index
         .columns
         .iter()
-        .position(|col| Some(&col.name) == column.name.as_ref())
+        .position(|col| Some(&col.name) == column.name.as_ref()))
 }
 
 /// Find all [IndexConstraint]s for a given WHERE clause
@@ -836,7 +880,7 @@ fn find_index_constraints(
                 continue;
             }
             // Skip terms that are not binary comparisons
-            let ast::Expr::Binary(lhs, operator, rhs) = &term.expr else {
+            let ast::Expr::Binary(lhs, operator, rhs) = unwrap_parens(&term.expr)? else {
                 continue;
             };
             // Only consider index scans for binary ops that are comparisons
@@ -851,9 +895,21 @@ fn find_index_constraints(
                 continue;
             }
 
+            // If both lhs and rhs refer to columns from this table, we can't use this constraint
+            // because we can't use the index to satisfy the condition.
+            // Examples:
+            // - WHERE t.x > t.y
+            // - WHERE t.x + 1 > t.y - 5
+            // - WHERE t.x = (t.x)
+            if determine_where_to_eval_expr(&lhs)? == EvalAt::Loop(table_index)
+                && determine_where_to_eval_expr(&rhs)? == EvalAt::Loop(table_index)
+            {
+                continue;
+            }
+
             // Check if lhs is a column that is in the i'th position of the index
             if Some(position_in_index)
-                == get_column_position_in_index(lhs, table_index, table_reference, index)
+                == get_column_position_in_index(lhs, table_index, table_reference, index)?
             {
                 out_constraints.push(IndexConstraint {
                     operator: *operator,
@@ -864,7 +920,7 @@ fn find_index_constraints(
             }
             // Check if rhs is a column that is in the i'th position of the index
             if Some(position_in_index)
-                == get_column_position_in_index(rhs, table_index, table_reference, index)
+                == get_column_position_in_index(rhs, table_index, table_reference, index)?
             {
                 out_constraints.push(IndexConstraint {
                     operator: opposite_cmp_op(*operator), // swap the operator since e.g. if condition is 5 >= x, we want to use x <= 5
@@ -916,7 +972,8 @@ pub fn build_seek_def_from_index_constraints(
         // Extract the other expression from the binary WhereTerm (i.e. the one being compared to the index column)
         let (idx, side) = constraint.position_in_where_clause;
         let where_term = &mut where_clause[idx];
-        let ast::Expr::Binary(lhs, _, rhs) = where_term.expr.take_ownership() else {
+        let ast::Expr::Binary(lhs, _, rhs) = unwrap_parens(where_term.expr.take_ownership())?
+        else {
             crate::bail_parse_error!("expected binary expression");
         };
         let cmp_expr = if side == BinaryExprSide::Lhs {
@@ -1186,6 +1243,16 @@ pub fn try_extract_rowid_search_expression(
     }
     match &mut cond.expr {
         ast::Expr::Binary(lhs, operator, rhs) => {
+            // If both lhs and rhs refer to columns from this table, we can't perform a rowid seek
+            // Examples:
+            // - WHERE t.x > t.y
+            // - WHERE t.x + 1 > t.y - 5
+            // - WHERE t.x = (t.x)
+            if determine_where_to_eval_expr(lhs)? == EvalAt::Loop(table_index)
+                && determine_where_to_eval_expr(rhs)? == EvalAt::Loop(table_index)
+            {
+                return Ok(None);
+            }
             if lhs.is_rowid_alias_of(table_index) {
                 match operator {
                     ast::Operator::Equals => {
diff --git a/core/translate/planner.rs b/core/translate/planner.rs
index b7b8745b..2d924666 100644
--- a/core/translate/planner.rs
+++ b/core/translate/planner.rs
@@ -570,7 +570,7 @@ pub fn parse_where(
   For expressions not referencing any tables (e.g. constants), this is before the main loop is
   opened, because they do not need any table data.
 */
-fn determine_where_to_eval_expr<'a>(predicate: &'a ast::Expr) -> Result<EvalAt> {
+pub fn determine_where_to_eval_expr<'a>(predicate: &'a ast::Expr) -> Result<EvalAt> {
     let mut eval_at: EvalAt = EvalAt::BeforeLoop;
     match predicate {
         ast::Expr::Binary(e1, _, e2) => {


Test patch:
diff --git a/testing/where.test b/testing/where.test
index a149e85f..958d7825 100755
--- a/testing/where.test
+++ b/testing/where.test
@@ -573,6 +573,12 @@ do_execsql_test where-constant-condition-no-tables-2 {
     select 1 where 1 IS NOT NULL;
 } {1}
 
+# We had a bug where NULL was incorrectly used as a seek key, returning all rows (because NULL < everything in index keys)
 do_execsql_test where-null-comparison-index-seek-regression-test {
     select age from users where age > NULL;
-} {}
\ No newline at end of file
+} {}
+
+# We had a bug where Limbo tried to use an index when there was a WHERE term like 't.x = t.x'
+do_execsql_test where-self-referential-regression {
+  select count(1) from users where id = id;
+} {10000}
diff --git a/tests/integration/fuzz/mod.rs b/tests/integration/fuzz/mod.rs
index 73263e5f..3814be97 100644
--- a/tests/integration/fuzz/mod.rs
+++ b/tests/integration/fuzz/mod.rs
@@ -1114,19 +1114,38 @@ mod tests {
         let limbo_conn = db.connect_limbo();
         let sqlite_conn = rusqlite::Connection::open_in_memory().unwrap();
         for table in tables.iter() {
-            let query = format!("CREATE TABLE {} ({})", table.name, table.columns.join(", "));
+            let columns_with_first_column_as_pk = {
+                let mut columns = vec![];
+                columns.push(format!("{} PRIMARY KEY", table.columns[0]));
+                columns.extend(table.columns[1..].iter().map(|c| c.to_string()));
+                columns.join(", ")
+            };
+            let query = format!(
+                "CREATE TABLE {} ({})",
+                table.name, columns_with_first_column_as_pk
+            );
+            let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
+            let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
+
             assert_eq!(
-                limbo_exec_rows(&db, &limbo_conn, &query),
-                sqlite_exec_rows(&sqlite_conn, &query)
+                limbo, sqlite,
+                "query: {}, limbo: {:?}, sqlite: {:?}",
+                query, limbo, sqlite
             );
         }
 
         let (mut rng, seed) = rng_from_time();
         log::info!("seed: {}", seed);
 
-        for _ in 0..100 {
-            let (x, y, z) = (
-                g.generate(&mut rng, builders.number, 1),
+        let mut i = 0;
+        let mut primary_key_set = HashSet::with_capacity(100);
+        while i < 100 {
+            let x = g.generate(&mut rng, builders.number, 1);
+            if primary_key_set.contains(&x) {
+                continue;
+            }
+            primary_key_set.insert(x.clone());
+            let (y, z) = (
                 g.generate(&mut rng, builders.number, 1),
                 g.generate(&mut rng, builders.number, 1),
             );
@@ -1138,7 +1157,13 @@ mod tests {
                 "seed: {}",
                 seed,
             );
+            i += 1;
         }
+        // verify the same number of rows in both tables
+        let query = format!("SELECT COUNT(*) FROM t");
+        let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
+        let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
+        assert_eq!(limbo, sqlite, "seed: {}", seed);
 
         let sql = g
             .create()
@@ -1152,11 +1177,25 @@ mod tests {
             log::info!("query: {}", query);
             let limbo = limbo_exec_rows(&db, &limbo_conn, &query);
             let sqlite = sqlite_exec_rows(&sqlite_conn, &query);
-            assert_eq!(
-                limbo, sqlite,
-                "query: {}, limbo: {:?}, sqlite: {:?} seed: {}",
-                query, limbo, sqlite, seed
-            );
+
+            if limbo.len() != sqlite.len() {
+                panic!("MISMATCHING ROW COUNT (limbo: {}, sqlite: {}) for query: {}\n\n limbo: {:?}\n\n sqlite: {:?}", limbo.len(), sqlite.len(), query, limbo, sqlite);
+            }
+            // find first row where limbo and sqlite differ
+            let diff_rows = limbo
+                .iter()
+                .zip(sqlite.iter())
+                .filter(|(l, s)| l != s)
+                .collect::<Vec<_>>();
+            if !diff_rows.is_empty() {
+                // due to different choices in index usage (usually in these cases sqlite is smart enough to use an index and we aren't),
+                // sqlite might return rows in a different order
+                // check if all limbo rows are present in sqlite
+                let all_present = limbo.iter().all(|l| sqlite.iter().any(|s| l == s));
+                if !all_present {
+                    panic!("MISMATCHING ROWS (limbo: {}, sqlite: {}) for query: {}\n\n limbo: {:?}\n\n sqlite: {:?}\n\n differences: {:?}", limbo.len(), sqlite.len(), query, limbo, sqlite, diff_rows);
+                }
+            }
         }
     }
 }