From be387eb31c2bf7534f7a738f7aa31074ae30fd57 Mon Sep 17 00:00:00 2001
From: Pere Diaz Bou <pere-altea@hotmail.com>
Date: Fri, 12 Jul 2024 17:00:05 +0200
Subject: [PATCH 1/3] core: labels system and limit refactor

Signed-off-by: Pere Diaz Bou <pere-altea@hotmail.com>
---
 core/translate.rs | 184 ++++++++++++++++++++++------------------------
 core/vdbe.rs      | 136 ++++++++++++++++++++++++++++++++--
 2 files changed, 216 insertions(+), 104 deletions(-)

diff --git a/core/translate.rs b/core/translate.rs
index 99a3c0a50..aaaeccc20 100644
--- a/core/translate.rs
+++ b/core/translate.rs
@@ -6,7 +6,7 @@ use crate::pager::Pager;
 use crate::schema::{Column, Schema, Table};
 use crate::sqlite3_ondisk::{DatabaseHeader, MIN_PAGE_CACHE_SIZE};
 use crate::util::normalize_ident;
-use crate::vdbe::{Insn, Program, ProgramBuilder};
+use crate::vdbe::{BranchOffset, Insn, Program, ProgramBuilder};
 use anyhow::Result;
 use sqlite3_parser::ast::{self, Expr};
 
@@ -31,7 +31,8 @@ struct Select {
 
 struct LoopInfo {
     table: Table,
-    rewind_offset: usize,
+    rewind_offset: BranchOffset,
+    rewind_label: BranchOffset,
     open_cursor: usize,
 }
 
@@ -60,6 +61,12 @@ impl ColumnInfo {
     }
 }
 
+struct LimitInfo {
+    limit_reg: usize,
+    num: i64,
+    goto_label: BranchOffset,
+}
+
 /// Translate SQL statement into bytecode program.
 pub fn translate(
     schema: &Schema,
@@ -161,59 +168,59 @@ fn build_select(schema: &Schema, select: ast::Select) -> Result<Select> {
 /// Generate code for a SELECT statement.
 fn translate_select(mut select: Select) -> Result<Program> {
     let mut program = ProgramBuilder::new();
-    let init_offset = program.emit_placeholder();
+    let init_label = program.allocate_label();
+    program.add_label(init_label, program.offset());
+    program.emit_insn(Insn::Init {
+        target_pc: init_label,
+    });
     let start_offset = program.offset();
-    let limit_reg = if let Some(limit) = &select.limit {
+
+    let limit_info = if let Some(limit) = &select.limit {
         assert!(limit.offset.is_none());
         let target_register = program.alloc_register();
-        Some(translate_expr(
-            &mut program,
-            &select,
-            &limit.expr,
-            target_register,
-        )?)
+        let limit_reg = translate_expr(&mut program, &select, &limit.expr, target_register)?;
+        let num = if let ast::Expr::Literal(ast::Literal::Numeric(num)) = &limit.expr {
+            num.parse::<i64>()?
+        } else {
+            todo!();
+        };
+        let goto_label = program.allocate_label();
+        if num == 0 {
+            program.add_label(goto_label, program.offset());
+            program.emit_insn(Insn::Goto {
+                target_pc: goto_label,
+            });
+        }
+        Some(LimitInfo {
+            limit_reg,
+            num,
+            goto_label,
+        })
     } else {
         None
     };
-    let parsed_limit = select.limit.as_ref().and_then(|limit| {
-        if let ast::Expr::Literal(ast::Literal::Numeric(num)) = &limit.expr {
-            num.parse::<i64>().ok()
-        } else {
-            None
-        }
-    });
-    let limit_goto = match parsed_limit {
-        Some(0) => Some(program.emit_placeholder()),
-        _ => None,
-    };
-    let limit_insn = if !select.src_tables.is_empty() {
+
+    if !select.src_tables.is_empty() {
         translate_tables_begin(&mut program, &mut select);
 
         let where_maybe = insert_where_clause_instructions(&select, &mut program)?;
 
         let (register_start, register_end) = translate_columns(&mut program, &select)?;
 
-        let mut limit_insn: Option<usize> = None;
         if !select.exist_aggregation {
             program.emit_insn(Insn::ResultRow {
                 start_reg: register_start,
                 count: register_end - register_start,
             });
-            limit_insn = limit_reg.map(|_| program.emit_placeholder());
+            emit_limit_insn(&limit_info, &mut program);
         }
 
-        let next_offset = translate_tables_end(&mut program, &select);
-
-        if let Some((where_clause_offset, where_reg)) = where_maybe {
-            program.fixup_insn(
-                where_clause_offset,
-                Insn::IfNot {
-                    reg: where_reg,
-                    target_pc: next_offset, // jump to 'next row' instruction if where not matched
-                },
-            );
+        if let Some((where_clause_label, where_reg)) = where_maybe {
+            program.resolve_label(where_clause_label, program.offset());
         }
 
+        translate_tables_end(&mut program, &select);
+
         if select.exist_aggregation {
             let mut target = register_start;
             for info in &select.column_info {
@@ -230,56 +237,27 @@ fn translate_select(mut select: Select) -> Result<Program> {
                 start_reg: register_start,
                 count: register_end - register_start,
             });
-            limit_insn = limit_reg.map(|_| program.emit_placeholder());
+            emit_limit_insn(&limit_info, &mut program);
         }
-        limit_insn
     } else {
         assert!(!select.exist_aggregation);
         let where_maybe = insert_where_clause_instructions(&select, &mut program)?;
         let (register_start, register_end) = translate_columns(&mut program, &select)?;
-        if let Some((where_clause_offset, where_reg)) = where_maybe {
-            program.fixup_insn(
-                where_clause_offset,
-                Insn::IfNot {
-                    reg: where_reg,
-                    target_pc: program.offset() + 1, // jump directly over the result row if no source tables and where not matched
-                },
-            );
+        if let Some((where_clause_label, where_reg)) = where_maybe {
+            program.resolve_label(where_clause_label, program.offset() + 1);
         }
         program.emit_insn(Insn::ResultRow {
             start_reg: register_start,
             count: register_end - register_start,
         });
-        limit_reg.map(|_| program.emit_placeholder())
+        emit_limit_insn(&limit_info, &mut program);
     };
     program.emit_insn(Insn::Halt);
     let halt_offset = program.offset() - 1;
-    if let Some(limit_goto) = limit_goto {
-        program.fixup_insn(
-            limit_goto,
-            Insn::Goto {
-                target_pc: halt_offset,
-            },
-        );
+    if limit_info.is_some() && limit_info.as_ref().unwrap().goto_label < 0 {
+        program.resolve_label(limit_info.as_ref().unwrap().goto_label, halt_offset);
     }
-    if let Some(limit_insn) = limit_insn {
-        let insn = match parsed_limit {
-            Some(0) => Insn::Goto {
-                target_pc: halt_offset,
-            },
-            _ => Insn::DecrJumpZero {
-                reg: limit_reg.unwrap(),
-                target_pc: halt_offset,
-            },
-        };
-        program.fixup_insn(limit_insn, insn);
-    }
-    program.fixup_insn(
-        init_offset,
-        Insn::Init {
-            target_pc: program.offset(),
-        },
-    );
+    program.resolve_label(init_label, program.offset());
     program.emit_insn(Insn::Transaction);
     program.emit_insn(Insn::Goto {
         target_pc: start_offset,
@@ -287,14 +265,34 @@ fn translate_select(mut select: Select) -> Result<Program> {
     Ok(program.build())
 }
 
+fn emit_limit_insn(limit_info: &Option<LimitInfo>, program: &mut ProgramBuilder) {
+    if limit_info.is_none() {
+        return;
+    }
+    let limit_info = limit_info.as_ref().unwrap();
+    if limit_info.num > 0 {
+        program.add_label(limit_info.goto_label, program.offset());
+        program.emit_insn(Insn::DecrJumpZero {
+            reg: limit_info.limit_reg,
+            target_pc: limit_info.goto_label,
+        });
+    }
+}
+
 fn insert_where_clause_instructions(
     select: &Select,
     program: &mut ProgramBuilder,
-) -> Result<Option<(usize, usize)>> {
+) -> Result<Option<(BranchOffset, usize)>> {
     if let Some(w) = &select.where_clause {
         let where_reg = program.alloc_register();
         let _ = translate_expr(program, &select, w, where_reg)?;
-        Ok(Some((program.emit_placeholder(), where_reg))) // We emit a placeholder because we determine the jump target later (after we know where the 'cursor next' instruction is)
+        let label = program.allocate_label();
+        program.add_label(label, program.offset());
+        program.emit_insn(Insn::IfNot {
+            reg: where_reg,
+            target_pc: label, // jump to 'next row' instruction if where not matched
+        });
+        Ok(Some((label, where_reg))) // We emit a placeholder because we determine the jump target later (after we know where the 'cursor next' instruction is)
     } else {
         Ok(None)
     }
@@ -312,29 +310,17 @@ fn translate_tables_begin(program: &mut ProgramBuilder, select: &mut Select) {
     }
 }
 
-fn translate_tables_end(program: &mut ProgramBuilder, select: &Select) -> usize {
+fn translate_tables_end(program: &mut ProgramBuilder, select: &Select) {
     // iterate in reverse order as we open cursors in order
-    let mut next_offset = std::usize::MAX;
     for table_loop in select.loops.iter().rev() {
         let cursor_id = table_loop.open_cursor;
         program.emit_insn(Insn::NextAsync { cursor_id });
-        if next_offset == std::usize::MAX {
-            next_offset = program.offset() - 1;
-        }
         program.emit_insn(Insn::NextAwait {
             cursor_id,
-            pc_if_next: table_loop.rewind_offset,
+            pc_if_next: table_loop.rewind_offset as BranchOffset,
         });
-        program.fixup_insn(
-            table_loop.rewind_offset,
-            Insn::RewindAwait {
-                cursor_id: table_loop.open_cursor,
-                pc_if_empty: program.offset(),
-            },
-        );
+        program.resolve_label(table_loop.rewind_label, program.offset());
     }
-
-    next_offset
 }
 
 fn translate_table_open_cursor(program: &mut ProgramBuilder, table: &Table) -> LoopInfo {
@@ -352,6 +338,7 @@ fn translate_table_open_cursor(program: &mut ProgramBuilder, table: &Table) -> L
         table: table.clone(),
         open_cursor: cursor_id,
         rewind_offset: 0,
+        rewind_label: 0,
     }
 }
 
@@ -359,8 +346,14 @@ fn translate_table_open_loop(program: &mut ProgramBuilder, loop_info: &mut LoopI
     program.emit_insn(Insn::RewindAsync {
         cursor_id: loop_info.open_cursor,
     });
-    let rewind_await_offset = program.emit_placeholder();
-    loop_info.rewind_offset = rewind_await_offset;
+    let rewind_await_label = program.allocate_label();
+    program.emit_insn(Insn::RewindAwait {
+        cursor_id: loop_info.open_cursor,
+        pc_if_empty: rewind_await_label,
+    });
+    program.add_label(rewind_await_label, program.offset() - 1);
+    loop_info.rewind_label = rewind_await_label;
+    loop_info.rewind_offset = program.offset() - 1;
 }
 
 fn translate_columns(program: &mut ProgramBuilder, select: &Select) -> Result<(usize, usize)> {
@@ -774,7 +767,11 @@ fn translate_pragma(
     pager: Rc<Pager>,
 ) -> Result<Program> {
     let mut program = ProgramBuilder::new();
-    let init_offset = program.emit_placeholder();
+    let init_label = program.allocate_label();
+    program.add_label(init_label, program.offset());
+    program.emit_insn(Insn::Init {
+        target_pc: init_label,
+    });
     let start_offset = program.offset();
     match body {
         None => {
@@ -811,12 +808,7 @@ fn translate_pragma(
         }
     };
     program.emit_insn(Insn::Halt);
-    program.fixup_insn(
-        init_offset,
-        Insn::Init {
-            target_pc: program.offset(),
-        },
-    );
+    program.resolve_label(init_label, program.offset());
     program.emit_insn(Insn::Transaction);
     program.emit_insn(Insn::Goto {
         target_pc: start_offset,
diff --git a/core/vdbe.rs b/core/vdbe.rs
index e99451d9d..57a11b716 100644
--- a/core/vdbe.rs
+++ b/core/vdbe.rs
@@ -6,15 +6,16 @@ use crate::types::{AggContext, Cursor, CursorResult, OwnedRecord, OwnedValue, Re
 use anyhow::Result;
 use std::borrow::BorrowMut;
 use std::cell::RefCell;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 use std::rc::Rc;
 
-pub type BranchOffset = usize;
+pub type BranchOffset = i64;
 
 pub type CursorID = usize;
 
 pub type PageIdx = usize;
 
+#[derive(Debug)]
 pub enum Insn {
     // Initialize the program state and jump to the given PC.
     Init {
@@ -205,18 +206,25 @@ pub enum Insn {
     },
 }
 
+// Index of insn in list of insns
+type InsnReference = usize;
+
 pub struct ProgramBuilder {
     next_free_register: usize,
+    next_free_label: BranchOffset,
     next_free_cursor_id: usize,
     insns: Vec<Insn>,
+    unresolved_labels: HashMap<BranchOffset, Vec<InsnReference>>,
 }
 
 impl ProgramBuilder {
     pub fn new() -> Self {
         Self {
             next_free_register: 0,
+            next_free_label: 0,
             next_free_cursor_id: 0,
             insns: Vec::new(),
+            unresolved_labels: HashMap::new(),
         }
     }
 
@@ -252,12 +260,124 @@ impl ProgramBuilder {
         self.insns.push(insn);
     }
 
-    pub fn fixup_insn(&mut self, offset: usize, insn: Insn) {
-        self.insns[offset] = insn;
+    pub fn offset(&self) -> BranchOffset {
+        self.insns.len() as BranchOffset
     }
 
-    pub fn offset(&self) -> usize {
-        self.insns.len()
+    pub fn allocate_label(&mut self) -> BranchOffset {
+        self.next_free_label -= 1;
+        self.unresolved_labels
+            .insert(self.next_free_label, Vec::new());
+        self.next_free_label
+    }
+
+    pub fn add_label(&mut self, label: BranchOffset, insn_reference: BranchOffset) {
+        assert!(insn_reference >= 0);
+        assert!(label < 0);
+        let insn_reference = insn_reference as InsnReference;
+        let label_references = self.unresolved_labels.get_mut(&label);
+        if let Some(label_references) = label_references {
+            label_references.push(insn_reference);
+        } else {
+            self.unresolved_labels.insert(label, vec![insn_reference]);
+        }
+    }
+
+    pub fn resolve_label(&mut self, label: BranchOffset, to_offset: BranchOffset) {
+        assert!(label < 0);
+        assert!(to_offset >= 0);
+        let label_references = self.unresolved_labels.get_mut(&label);
+        if label_references.is_none() {
+            unreachable!("Forbidden resolve of an unexistent label!");
+        }
+        let label_references = label_references.unwrap();
+
+        for insn_reference in label_references {
+            let insn = &mut self.insns[*insn_reference];
+            match insn {
+                Insn::Init { target_pc } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Eq {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Ne {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Lt {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Le {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Gt {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::Ge {
+                    lhs,
+                    rhs,
+                    target_pc,
+                } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::IfNot { reg, target_pc } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::RewindAwait {
+                    cursor_id,
+                    pc_if_empty,
+                } => {
+                    assert!(*pc_if_empty < 0);
+                    *pc_if_empty = to_offset;
+                }
+                Insn::Goto { target_pc } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::DecrJumpZero { reg, target_pc } => {
+                    assert!(*target_pc < 0);
+                    *target_pc = to_offset;
+                }
+                Insn::SorterNext {
+                    cursor_id,
+                    pc_if_next,
+                } => {
+                    assert!(*pc_if_next < 0);
+                    *pc_if_next = to_offset;
+                }
+                _ => {
+                    todo!("missing resolve_label for {:?}", insn);
+                }
+            }
+        }
     }
 
     pub fn build(self) -> Program {
@@ -316,7 +436,7 @@ impl Program {
         let mut prev_insn: Option<&Insn> = None;
         for (addr, insn) in self.insns.iter().enumerate() {
             indent_count = get_indent_count(indent_count, insn, prev_insn);
-            print_insn(addr, insn, indent.repeat(indent_count));
+            print_insn(addr as BranchOffset, insn, indent.repeat(indent_count));
             prev_insn = Some(insn);
         }
     }
@@ -327,7 +447,7 @@ impl Program {
         pager: Rc<Pager>,
     ) -> Result<StepResult<'a>> {
         loop {
-            let insn = &self.insns[state.pc];
+            let insn = &self.insns[state.pc as usize];
             trace_insn(state.pc, insn);
             let mut cursors = state.cursors.borrow_mut();
             match insn {

From d0f37969cee4ba802ce13f2b8e2c355b30725271 Mon Sep 17 00:00:00 2001
From: Pere Diaz Bou <pere-altea@hotmail.com>
Date: Sat, 13 Jul 2024 18:53:35 +0200
Subject: [PATCH 2/3] core: matrix of labelsXinsnref instead of hashmap

---
 core/vdbe.rs | 38 +++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/core/vdbe.rs b/core/vdbe.rs
index 57a11b716..e57f9ac60 100644
--- a/core/vdbe.rs
+++ b/core/vdbe.rs
@@ -214,7 +214,9 @@ pub struct ProgramBuilder {
     next_free_label: BranchOffset,
     next_free_cursor_id: usize,
     insns: Vec<Insn>,
-    unresolved_labels: HashMap<BranchOffset, Vec<InsnReference>>,
+    // Each lable has a list of InsnRefereces that must
+    // be resolved. Lists are indexed by: label.abs() - 1
+    unresolved_labels: Vec<Vec<InsnReference>>,
 }
 
 impl ProgramBuilder {
@@ -224,7 +226,7 @@ impl ProgramBuilder {
             next_free_label: 0,
             next_free_cursor_id: 0,
             insns: Vec::new(),
-            unresolved_labels: HashMap::new(),
+            unresolved_labels: Vec::new(),
         }
     }
 
@@ -266,32 +268,38 @@ impl ProgramBuilder {
 
     pub fn allocate_label(&mut self) -> BranchOffset {
         self.next_free_label -= 1;
-        self.unresolved_labels
-            .insert(self.next_free_label, Vec::new());
+        self.unresolved_labels.push(Vec::new());
         self.next_free_label
     }
 
+    fn label_to_index(&self, label: BranchOffset) -> usize {
+        (label.abs() - 1) as usize
+    }
+
     pub fn add_label(&mut self, label: BranchOffset, insn_reference: BranchOffset) {
         assert!(insn_reference >= 0);
         assert!(label < 0);
+        let label_index = self.label_to_index(label);
+        assert!(label_index < self.unresolved_labels.len());
         let insn_reference = insn_reference as InsnReference;
-        let label_references = self.unresolved_labels.get_mut(&label);
-        if let Some(label_references) = label_references {
-            label_references.push(insn_reference);
-        } else {
-            self.unresolved_labels.insert(label, vec![insn_reference]);
-        }
+        let label_references = &mut self.unresolved_labels[label_index];
+        label_references.push(insn_reference);
     }
 
     pub fn resolve_label(&mut self, label: BranchOffset, to_offset: BranchOffset) {
         assert!(label < 0);
         assert!(to_offset >= 0);
-        let label_references = self.unresolved_labels.get_mut(&label);
-        if label_references.is_none() {
-            unreachable!("Forbidden resolve of an unexistent label!");
-        }
-        let label_references = label_references.unwrap();
+        let label_index = self.label_to_index(label);
+        assert!(
+            label_index < self.unresolved_labels.len(),
+            "Forbidden resolve of an unexistent label!"
+        );
 
+        let label_references = &self.unresolved_labels[self.label_to_index(label)];
+        assert!(
+            label_references.len() > 0,
+            "Trying to resolve an empty created label, all labels must be resolved for now."
+        );
         for insn_reference in label_references {
             let insn = &mut self.insns[*insn_reference];
             match insn {

From 76010451221ec179e9632fb456cf8eff5b48d4f8 Mon Sep 17 00:00:00 2001
From: Pere Diaz Bou <pere-altea@hotmail.com>
Date: Sat, 13 Jul 2024 19:04:00 +0200
Subject: [PATCH 3/3] core: clear references on resolve

---
 core/vdbe.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/core/vdbe.rs b/core/vdbe.rs
index e57f9ac60..33e09dd6e 100644
--- a/core/vdbe.rs
+++ b/core/vdbe.rs
@@ -295,12 +295,12 @@ impl ProgramBuilder {
             "Forbidden resolve of an unexistent label!"
         );
 
-        let label_references = &self.unresolved_labels[self.label_to_index(label)];
+        let label_references = &mut self.unresolved_labels[label_index];
         assert!(
             label_references.len() > 0,
             "Trying to resolve an empty created label, all labels must be resolved for now."
         );
-        for insn_reference in label_references {
+        for insn_reference in label_references.iter() {
             let insn = &mut self.insns[*insn_reference];
             match insn {
                 Insn::Init { target_pc } => {
@@ -386,6 +386,7 @@ impl ProgramBuilder {
                 }
             }
         }
+        label_references.clear();
     }
 
     pub fn build(self) -> Program {