1use crate::{alloc, turso_assert, turso_assert_eq, turso_debug_assert, Result};
2
3use rustc_hash::FxHashMap as HashMap;
4use tracing::{instrument, Level};
5use turso_parser::ast::{self, ResolveType, SortOrder, TableInternalId};
6
7use crate::{
8 index_method::IndexMethodAttachment,
9 parameters::Parameters,
10 schema::{BTreeTable, Column, ColumnLayout, Index, PseudoCursorType, Schema, Table, Trigger},
11 translate::{
12 collate::CollationSeq,
13 emitter::{MaterializedColumnRef, TransactionMode},
14 plan::{ResultSetColumn, TableReferences},
15 },
16 Arc, CaptureDataChangesInfo, Connection, VirtualTable,
17};
18
19const HASH_TABLE_ID_BASE: usize = 1 << 30;
21
22#[derive(Default)]
23pub struct TableRefIdCounter {
24 next_free: ast::TableInternalId,
25}
26
27impl TableRefIdCounter {
28 pub fn new() -> Self {
29 Self {
30 next_free: TableInternalId::default(),
31 }
32 }
33
34 #[allow(clippy::should_implement_trait)]
35 pub fn next(&mut self) -> ast::TableInternalId {
36 let id = self.next_free;
37 self.next_free += 1;
38 id
39 }
40}
41
42use super::{
43 affinity::Affinity, BranchOffset, CursorID, Insn, InsnReference, PrepareContext,
44 PreparedProgram, Program,
45};
46use crate::translate::plan::BitSet;
47use std::num::NonZeroUsize;
48
49#[derive(Debug, Clone)]
53pub struct CursorKey {
54 pub table_reference_id: TableInternalId,
60 pub index: Option<Arc<Index>>,
63 pub is_build: bool,
65}
66
67impl CursorKey {
68 pub fn table(table_reference_id: TableInternalId) -> Self {
69 Self {
70 table_reference_id,
71 index: None,
72 is_build: false,
73 }
74 }
75
76 pub fn index(table_reference_id: TableInternalId, index: Arc<Index>) -> Self {
77 Self {
78 table_reference_id,
79 index: Some(index),
80 is_build: false,
81 }
82 }
83
84 pub fn hash_build(table_reference_id: TableInternalId) -> Self {
87 Self {
88 table_reference_id,
89 index: None,
90 is_build: true,
91 }
92 }
93
94 pub fn equals(&self, other: &CursorKey) -> bool {
95 if self.table_reference_id != other.table_reference_id {
96 return false;
97 }
98 if self.is_build != other.is_build {
99 return false;
100 }
101 match (self.index.as_ref(), other.index.as_ref()) {
102 (Some(self_index), Some(other_index)) => self_index.name == other_index.name,
103 (None, None) => true,
104 _ => false,
105 }
106 }
107}
108
109#[derive(Clone)]
111pub enum SelfTableContext {
112 ForSelect {
113 table_ref_id: TableInternalId,
114 referenced_tables: TableReferences,
115 },
116 ForDML {
117 dml_ctx: DmlColumnContext,
118 table: Arc<BTreeTable>,
119 },
120}
121
122#[derive(Clone)]
123enum DmlColumnRegisters {
124 Layout {
126 base_reg: usize,
127 rowid_reg: usize,
128 layout: ColumnLayout,
129 },
130 Indexed {
131 column_regs: Vec<usize>,
132 },
133}
134
135#[derive(Clone)]
136pub struct DmlColumnContext {
137 registers: DmlColumnRegisters,
138 rowid_alias_col: Option<usize>,
139}
140
141impl DmlColumnContext {
142 pub fn layout(
143 columns: &[Column],
144 base_reg: usize,
145 rowid_reg: usize,
146 layout: ColumnLayout,
147 ) -> Self {
148 let rowid_alias_col = columns.iter().position(|c| c.is_rowid_alias());
149
150 Self {
151 registers: DmlColumnRegisters::Layout {
152 base_reg,
153 rowid_reg,
154 layout,
155 },
156 rowid_alias_col,
157 }
158 }
159
160 pub fn from_column_reg_mapping<'a>(pairs: impl Iterator<Item = (&'a Column, usize)>) -> Self {
161 let mut rowid_alias_col = None;
162 let mut column_regs = Vec::new();
163 for (idx, (col, reg)) in pairs.enumerate() {
164 column_regs.push(reg);
165 if col.is_rowid_alias() {
166 rowid_alias_col = Some(idx);
167 }
168 }
169 Self {
170 registers: DmlColumnRegisters::Indexed { column_regs },
171 rowid_alias_col,
172 }
173 }
174
175 pub fn to_column_reg(&self, col_idx: usize) -> usize {
176 match &self.registers {
177 DmlColumnRegisters::Layout {
178 base_reg,
179 rowid_reg,
180 layout,
181 } => {
182 if self.rowid_alias_col == Some(col_idx) {
183 *rowid_reg
184 } else {
185 layout.to_register(*base_reg, col_idx)
186 }
187 }
188 DmlColumnRegisters::Indexed { column_regs } => column_regs[col_idx],
189 }
190 }
191}
192
193pub struct ProgramBuilder {
194 pub constant_spans: Vec<(usize, usize)>,
198 pub cursor_ref: Vec<(Option<CursorKey>, CursorType)>,
203 label_to_resolved_offset: Vec<Option<InsnReference>>,
209 comments: Vec<(InsnReference, &'static str)>,
211 pub parameters: Parameters,
212 pub result_columns: Vec<ResultSetColumn>,
213 pub insns: Vec<(Insn, usize)>,
215 materialized_ctes: HashMap<usize, MaterializedCteInfo>,
218 ctes_being_defined: Vec<String>,
222 pub trigger: Option<Arc<Trigger>>,
224 pub table_reference_counter: TableRefIdCounter,
225 collation: Option<(CollationSeq, bool)>,
227 capture_data_changes_info: Option<CaptureDataChangesInfo>,
228 mvcc_enabled: bool,
230 txn_mode: TransactionMode,
232 write_databases: BitSet,
234 read_databases: BitSet,
236 write_database_cookies: HashMap<usize, u32>,
238 read_database_cookies: HashMap<usize, u32>,
240 cursor_overrides: HashMap<usize, CursorID>,
243 pub id_register_overrides: HashMap<String, usize>,
247 hash_build_signatures: HashMap<usize, HashBuildSignature>,
249 hash_tables_to_keep_open: BitSet,
251 subquery_result_regs: HashMap<TableInternalId, usize>,
254 query_mode: QueryMode,
256 pub flags: ProgramBuilderFlags,
257 emitted_function_call: bool,
259 next_free_register: usize,
260 next_free_cursor_id: usize,
261 next_hash_table_id: usize,
262 pub table_references: TableReferences,
263 nested_level: usize,
265 init_label: BranchOffset,
266 start_offset: BranchOffset,
267 current_parent_explain_idx: Option<usize>,
269 pub(crate) reg_result_cols_start: Option<usize>,
270 pub resolve_type: ResolveType,
271 pub trigger_conflict_override: Option<ResolveType>,
275 next_cte_id: usize,
278 next_subquery_eqp_id: usize,
280 pub(crate) target_union_type: Option<Arc<crate::schema::TypeDef>>,
295}
296
297#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
298#[repr(transparent)]
299pub struct ProgramBuilderFlags(u8);
300
301impl ProgramBuilderFlags {
302 const ROLLBACK: u8 = 1 << 0;
303 const IS_MULTI_WRITE: u8 = 1 << 1;
304 const MAY_ABORT: u8 = 1 << 2;
305 const READONLY: u8 = 1 << 3;
306 const IS_SUBPROGRAM: u8 = 1 << 4;
307 const HAS_STATEMENT_CONFLICT: u8 = 1 << 5;
308 const SUPPRESS_CUSTOM_TYPE_DECODE: u8 = 1 << 6;
309 const SUPPRESS_COLUMN_DEFAULT: u8 = 1 << 7;
310
311 const fn new(is_subprogram: bool) -> Self {
312 let mut new = Self(0);
313 new.set_is_multi_write(true);
314 new.set_may_abort(true);
315 new.set_readonly(true);
316 new.set_is_subprogram(is_subprogram);
317 new.set_is_multi_write(true);
318 new.set_may_abort(true);
319 new
320 }
321
322 #[inline]
323 const fn get(self, bit: u8) -> bool {
324 (self.0 & bit) != 0
325 }
326
327 #[inline]
328 const fn set(&mut self, bit: u8, value: bool) {
329 if value {
330 self.0 |= bit;
331 } else {
332 self.0 &= !bit;
333 }
334 }
335
336 #[inline]
337 pub const fn rollback(self) -> bool {
338 self.get(Self::ROLLBACK)
339 }
340 #[inline]
341 pub const fn set_rollback(&mut self, v: bool) {
342 self.set(Self::ROLLBACK, v)
343 }
344
345 #[inline]
346 pub const fn is_multi_write(self) -> bool {
352 self.get(Self::IS_MULTI_WRITE)
353 }
354 #[inline]
355 pub const fn set_is_multi_write(&mut self, v: bool) {
356 self.set(Self::IS_MULTI_WRITE, v)
357 }
358
359 #[inline]
360 pub const fn may_abort(self) -> bool {
364 self.get(Self::MAY_ABORT)
365 }
366 #[inline]
367 pub const fn set_may_abort(&mut self, v: bool) {
368 self.set(Self::MAY_ABORT, v)
369 }
370
371 #[inline]
372 pub const fn readonly(self) -> bool {
376 self.get(Self::READONLY)
377 }
378 #[inline]
379 pub const fn set_readonly(&mut self, v: bool) {
380 self.set(Self::READONLY, v)
381 }
382
383 #[inline]
384 pub const fn is_subprogram(self) -> bool {
386 self.get(Self::IS_SUBPROGRAM)
387 }
388 #[inline]
389 pub const fn set_is_subprogram(&mut self, v: bool) {
390 self.set(Self::IS_SUBPROGRAM, v)
391 }
392
393 #[inline]
394 pub fn has_statement_conflict(self) -> bool {
397 self.get(Self::HAS_STATEMENT_CONFLICT)
398 }
399 #[inline]
400 pub fn set_has_statement_conflict(&mut self, v: bool) {
401 self.set(Self::HAS_STATEMENT_CONFLICT, v)
402 }
403
404 #[inline]
405 pub const fn suppress_custom_type_decode(self) -> bool {
409 self.get(Self::SUPPRESS_CUSTOM_TYPE_DECODE)
410 }
411 #[inline]
412 pub const fn set_suppress_custom_type_decode(&mut self, v: bool) {
413 self.set(Self::SUPPRESS_CUSTOM_TYPE_DECODE, v)
414 }
415
416 #[inline]
417 pub const fn suppress_column_default(self) -> bool {
421 self.get(Self::SUPPRESS_COLUMN_DEFAULT)
422 }
423 #[inline]
424 pub const fn set_suppress_column_default(&mut self, v: bool) {
425 self.set(Self::SUPPRESS_COLUMN_DEFAULT, v)
426 }
427}
428
429#[derive(Debug, Clone, Copy, Eq, PartialEq)]
430pub enum MaterializedBuildInputModeTag {
431 RowidOnly,
432 Payload,
433}
434
435#[derive(Debug, Clone, PartialEq, Eq)]
436pub struct HashBuildSignature {
439 pub join_key_indices: Vec<usize>,
441 pub payload_refs: Vec<MaterializedColumnRef>,
443 pub key_affinities: String,
445 pub use_bloom_filter: bool,
447 pub materialized_input_cursor: Option<CursorID>,
449 pub materialized_mode: Option<MaterializedBuildInputModeTag>,
451}
452
453#[derive(Debug, Clone)]
455pub struct MaterializedCteInfo {
456 pub cursor_id: CursorID,
458 pub table: Arc<BTreeTable>,
460 pub num_columns: usize,
462}
463
464#[derive(Debug, Clone)]
465pub enum CursorType {
466 BTreeTable(Arc<BTreeTable>),
467 BTreeIndex(Arc<Index>),
468 IndexMethod(Arc<dyn IndexMethodAttachment>),
469 Pseudo(PseudoCursorType),
470 Sorter,
471 VirtualTable(Arc<VirtualTable>),
472 MaterializedView(
473 Arc<BTreeTable>,
474 Arc<crate::sync::Mutex<crate::incremental::view::IncrementalView>>,
475 ),
476}
477
478impl CursorType {
479 pub const fn is_index(&self) -> bool {
480 matches!(self, CursorType::BTreeIndex(_))
481 }
482
483 pub fn get_explain_description(&self) -> String {
484 let out = match self {
485 CursorType::BTreeTable(btree_table) => {
486 let mut col_count = btree_table.columns().len();
487 if btree_table.get_rowid_alias_column().is_none() {
488 col_count += 1;
489 }
490 Some((
491 col_count,
492 btree_table
493 .columns()
494 .iter()
495 .map(|col| {
496 if let Some(coll) = col.collation_opt() {
497 format!("{coll}")
498 } else {
499 "B".to_string()
500 }
501 })
502 .collect::<Vec<_>>()
503 .join(","),
504 ))
505 }
506 CursorType::BTreeIndex(index) => {
507 let mut col_count = index.columns.len();
508 if index.has_rowid {
509 col_count += 1;
510 }
511 Some((
512 col_count,
513 index
514 .columns
515 .iter()
516 .map(|col| {
517 let sign = match col.order {
518 SortOrder::Asc => "",
519 SortOrder::Desc => "-",
520 };
521 if let Some(coll) = col.collation {
522 format!("{sign}{coll}")
523 } else {
524 format!("{sign}B")
525 }
526 })
527 .collect::<Vec<_>>()
528 .join(","),
529 ))
530 }
531 _ => None,
532 };
533
534 out.map_or(String::new(), |(col_count, collations)| {
535 format!("k({col_count},{collations})")
536 })
537 }
538}
539
540#[derive(Debug, Clone, PartialEq, Eq, Copy)]
541pub enum QueryMode {
542 Normal,
543 Explain,
544 ExplainQueryPlan,
545}
546
547impl QueryMode {
548 pub const fn new(cmd: &ast::Cmd) -> Self {
549 match cmd {
550 ast::Cmd::ExplainQueryPlan(_) => QueryMode::ExplainQueryPlan,
551 ast::Cmd::Explain(_) => QueryMode::Explain,
552 ast::Cmd::Stmt(_) => QueryMode::Normal,
553 }
554 }
555}
556
557pub struct ProgramBuilderOpts {
558 pub num_cursors: usize,
559 pub approx_num_insns: usize,
560 pub approx_num_labels: usize,
561}
562
563impl ProgramBuilderOpts {
564 pub const fn new(
565 num_cursors: usize,
566 approx_num_insns: usize,
567 approx_num_labels: usize,
568 ) -> Self {
569 Self {
570 num_cursors,
571 approx_num_insns,
572 approx_num_labels,
573 }
574 }
575}
576
577#[macro_export]
581macro_rules! emit_explain {
582 ($builder:expr, $push:expr, $detail:expr) => {
583 if let $crate::QueryMode::ExplainQueryPlan = $builder.get_query_mode() {
584 $builder.emit_explain($push, $detail);
585 }
586 };
587}
588
589impl ProgramBuilder {
590 pub fn register_variable(&mut self, variable: &ast::Variable) -> NonZeroUsize {
593 let index = usize::try_from(variable.index.get())
594 .expect("u32 variable index must fit into usize")
595 .try_into()
596 .expect("variable index must be non-zero");
597 if let Some(name) = variable.name.as_deref() {
598 self.parameters.push_named_at(name, index);
599 } else {
600 self.parameters.push_index(index);
601 }
602 index
603 }
604
605 pub fn with_scoped_result_cols_start<T>(
608 &mut self,
609 f: impl FnOnce(&mut Self) -> crate::Result<T>,
610 ) -> crate::Result<T> {
611 let saved = self.reg_result_cols_start;
612 let result = f(self);
613 self.reg_result_cols_start = saved;
614 result
615 }
616
617 pub fn new(
618 query_mode: QueryMode,
619 capture_data_changes_info: Option<CaptureDataChangesInfo>,
620 opts: ProgramBuilderOpts,
621 ) -> Self {
622 ProgramBuilder::_new(query_mode, capture_data_changes_info, opts, None, false)
623 }
624 pub fn new_for_trigger(
625 query_mode: QueryMode,
626 capture_data_changes_info: Option<CaptureDataChangesInfo>,
627 opts: ProgramBuilderOpts,
628 trigger: Arc<Trigger>,
629 ) -> Self {
630 ProgramBuilder::_new(
631 query_mode,
632 capture_data_changes_info,
633 opts,
634 Some(trigger),
635 true,
636 )
637 }
638 pub fn new_for_subprogram(
641 query_mode: QueryMode,
642 capture_data_changes_info: Option<CaptureDataChangesInfo>,
643 opts: ProgramBuilderOpts,
644 ) -> Self {
645 ProgramBuilder::_new(query_mode, capture_data_changes_info, opts, None, true)
646 }
647
648 #[turso_macros::trace_stack]
649 fn _new(
650 query_mode: QueryMode,
651 capture_data_changes_info: Option<CaptureDataChangesInfo>,
652 opts: ProgramBuilderOpts,
653 trigger: Option<Arc<Trigger>>,
654 is_subprogram: bool,
655 ) -> Self {
656 Self {
657 table_reference_counter: TableRefIdCounter::new(),
658 next_free_register: 1,
659 next_free_cursor_id: 0,
660 next_hash_table_id: HASH_TABLE_ID_BASE,
661 insns: Vec::with_capacity(opts.approx_num_insns),
662 cursor_ref: Vec::with_capacity(opts.num_cursors),
663 constant_spans: Vec::new(),
664 label_to_resolved_offset: Vec::with_capacity(opts.approx_num_labels),
665 comments: Vec::new(),
666 parameters: Parameters::new(),
667 result_columns: Vec::new(),
668 table_references: TableReferences::new(vec![], vec![]),
669 collation: None,
670 nested_level: 0,
671 init_label: BranchOffset::Placeholder,
673 start_offset: BranchOffset::Placeholder,
674 capture_data_changes_info,
675 mvcc_enabled: false,
676 txn_mode: TransactionMode::None,
677 write_databases: BitSet::default(),
678 read_databases: BitSet::default(),
679 write_database_cookies: HashMap::default(),
680 read_database_cookies: HashMap::default(),
681 query_mode,
682 current_parent_explain_idx: None,
683 reg_result_cols_start: None,
684 flags: ProgramBuilderFlags::new(is_subprogram),
685 emitted_function_call: false,
686 trigger,
687 resolve_type: ResolveType::Abort,
688 trigger_conflict_override: None,
689 cursor_overrides: HashMap::default(),
690 id_register_overrides: HashMap::default(),
691 hash_build_signatures: HashMap::default(),
692 hash_tables_to_keep_open: BitSet::default(),
693 subquery_result_regs: HashMap::default(),
694 next_cte_id: 0,
695 materialized_ctes: HashMap::default(),
696 ctes_being_defined: Vec::new(),
697 next_subquery_eqp_id: 1,
698 target_union_type: None,
699 }
700 }
701
702 pub const fn next_subquery_eqp_id(&mut self) -> usize {
703 let id = self.next_subquery_eqp_id;
704 self.next_subquery_eqp_id += 1;
705 id
706 }
707
708 pub const fn alloc_hash_table_id(&mut self) -> usize {
709 let id = self.next_hash_table_id;
710 self.next_hash_table_id = self
711 .next_hash_table_id
712 .checked_add(1)
713 .expect("hash table id overflow");
714 id
715 }
716
717 pub const fn alloc_cte_id(&mut self) -> usize {
720 let id = self.next_cte_id;
721 self.next_cte_id += 1;
722 id
723 }
724
725 pub fn get_materialized_cte(&self, cte_id: usize) -> Option<&MaterializedCteInfo> {
728 self.materialized_ctes.get(&cte_id)
729 }
730
731 pub fn register_materialized_cte(&mut self, cte_id: usize, info: MaterializedCteInfo) {
733 self.materialized_ctes.insert(cte_id, info);
734 }
735
736 pub fn push_cte_being_defined(&mut self, name: String) {
740 self.ctes_being_defined.push(name);
741 }
742
743 pub fn pop_cte_being_defined(&mut self) {
745 self.ctes_being_defined.pop();
746 }
747
748 pub fn is_cte_being_defined(&self, name: &str) -> bool {
750 self.ctes_being_defined.iter().any(|n| n == name)
751 }
752
753 pub fn take_ctes_being_defined(&mut self) -> Vec<String> {
756 std::mem::take(&mut self.ctes_being_defined)
757 }
758
759 pub fn restore_ctes_being_defined(&mut self, saved: Vec<String>) {
761 self.ctes_being_defined = saved;
762 }
763
764 pub const fn set_resolve_type(&mut self, resolve_type: ResolveType) {
765 self.resolve_type = resolve_type;
766 }
767
768 pub const fn set_trigger_conflict_override(&mut self, resolve_type: ResolveType) {
771 self.trigger_conflict_override = Some(resolve_type);
772 }
773
774 pub fn should_keep_hash_table_open(&self, hash_table_id: usize) -> bool {
776 self.hash_tables_to_keep_open.get(hash_table_id)
777 }
778
779 pub fn set_hash_tables_to_keep_open(&mut self, tables: &BitSet) {
781 self.hash_tables_to_keep_open.clone_from(tables);
782 }
783
784 pub fn clear_hash_tables_to_keep_open(&mut self) {
786 self.hash_tables_to_keep_open = BitSet::default();
787 }
788
789 pub fn hash_build_signature_matches(
791 &self,
792 hash_table_id: usize,
793 signature: &HashBuildSignature,
794 ) -> bool {
795 self.hash_build_signatures
796 .get(&hash_table_id)
797 .is_some_and(|existing| existing == signature)
798 }
799
800 pub fn has_hash_build_signature(&self, hash_table_id: usize) -> bool {
802 self.hash_build_signatures.contains_key(&hash_table_id)
803 }
804
805 pub fn record_hash_build_signature(
807 &mut self,
808 hash_table_id: usize,
809 signature: HashBuildSignature,
810 ) {
811 self.hash_build_signatures.insert(hash_table_id, signature);
812 }
813
814 pub fn clear_hash_build_signature(&mut self, hash_table_id: usize) {
816 self.hash_build_signatures.remove(&hash_table_id);
817 }
818
819 pub fn set_subquery_result_reg(&mut self, internal_id: TableInternalId, result_reg: usize) {
822 self.subquery_result_regs.insert(internal_id, result_reg);
823 }
824
825 pub fn get_subquery_result_reg(&self, internal_id: TableInternalId) -> Option<usize> {
828 self.subquery_result_regs.get(&internal_id).copied()
829 }
830
831 pub const fn set_multi_write(&mut self, is_multi_write: bool) {
834 self.flags.set_is_multi_write(is_multi_write);
835 }
836
837 pub const fn set_may_abort(&mut self, may_abort: bool) {
839 self.flags.set_may_abort(may_abort);
840 }
841
842 pub const fn may_abort(&self) -> bool {
849 self.flags.may_abort() || self.emitted_function_call
850 }
851
852 pub const fn capture_data_changes_info(&self) -> &Option<CaptureDataChangesInfo> {
853 &self.capture_data_changes_info
854 }
855
856 pub const fn is_mvcc_enabled(&self) -> bool {
858 self.mvcc_enabled
859 }
860
861 pub fn set_mvcc_enabled(&mut self, enabled: bool) {
862 self.mvcc_enabled = enabled;
863 }
864
865 pub fn extend(&mut self, opts: &ProgramBuilderOpts) {
866 self.insns.reserve(opts.approx_num_insns);
867 self.cursor_ref.reserve(opts.num_cursors);
868 self.label_to_resolved_offset
869 .reserve(opts.approx_num_labels);
870 }
871
872 pub fn constant_span_start(&mut self) -> usize {
875 let span = self.constant_spans.len();
876 let start = self.insns.len();
877 self.constant_spans.push((start, usize::MAX));
878 span
879 }
880
881 pub fn constant_span_end(&mut self, span_idx: usize) {
884 let span = &mut self.constant_spans[span_idx];
885 if span.1 == usize::MAX {
886 span.1 = self.insns.len().saturating_sub(1);
887 }
888 }
889
890 pub fn constant_span_end_all(&mut self) {
894 for span in self.constant_spans.iter_mut() {
895 if span.1 == usize::MAX {
896 span.1 = self.insns.len().saturating_sub(1);
897 }
898 }
899 }
900
901 pub fn constant_span_is_open(&self) -> bool {
903 self.constant_spans
904 .last()
905 .is_some_and(|(_, end)| *end == usize::MAX)
906 }
907
908 pub const fn constant_spans_next_idx(&self) -> usize {
912 self.constant_spans.len()
913 }
914
915 pub fn constant_spans_invalidate_after(&mut self, idx: usize) {
919 self.constant_spans.truncate(idx);
920 }
921
922 pub const fn alloc_register(&mut self) -> usize {
923 let reg = self.next_free_register;
924 self.next_free_register += 1;
925 reg
926 }
927
928 pub const fn alloc_registers(&mut self, amount: usize) -> usize {
929 let reg = self.next_free_register;
930 self.next_free_register += amount;
931 reg
932 }
933
934 pub const fn peek_next_register(&self) -> usize {
936 self.next_free_register
937 }
938
939 pub fn alloc_registers_and_init_w_null(&mut self, amount: usize) -> usize {
940 let reg = self.alloc_registers(amount);
941 self.emit_insn(Insn::Null {
942 dest: reg,
943 dest_end: if amount == 1 {
944 None
945 } else {
946 Some(reg + amount - 1)
947 },
948 });
949 reg
950 }
951
952 pub fn alloc_cursor_id_keyed(&mut self, key: CursorKey, cursor_type: CursorType) -> usize {
953 turso_assert!(
954 !self
955 .cursor_ref
956 .iter()
957 .any(|(k, _)| k.as_ref().is_some_and(|k| k.equals(&key))),
958 "duplicate cursor key"
959 );
960 self._alloc_cursor_id(Some(key), cursor_type)
961 }
962
963 pub fn alloc_cursor_id_keyed_if_not_exists(
964 &mut self,
965 key: CursorKey,
966 cursor_type: CursorType,
967 ) -> usize {
968 if let Some(cursor_id) = self.resolve_cursor_id_safe(&key) {
969 cursor_id
970 } else {
971 self._alloc_cursor_id(Some(key), cursor_type)
972 }
973 }
974
975 pub fn alloc_cursor_index(
977 &mut self,
978 key: Option<CursorKey>,
979 index: &Arc<Index>,
980 ) -> crate::Result<usize> {
981 tracing::debug!("alloc cursor: {:?} {:?}", key, index.index_method.is_some());
982 let module = index.index_method.as_ref();
983 if let Some(m) = module {
984 if !m.definition().backing_btree {
985 return Ok(self._alloc_cursor_id(key, CursorType::IndexMethod(m.clone())));
986 }
987 }
988 Ok(self._alloc_cursor_id(key, CursorType::BTreeIndex(index.clone())))
989 }
990
991 pub fn alloc_cursor_index_if_not_exists(
992 &mut self,
993 key: CursorKey,
994 index: &Arc<Index>,
995 ) -> crate::Result<usize> {
996 if let Some(cursor_id) = self.resolve_cursor_id_safe(&key) {
997 Ok(cursor_id)
998 } else {
999 self.alloc_cursor_index(Some(key), index)
1000 }
1001 }
1002
1003 pub fn alloc_cursor_id(&mut self, cursor_type: CursorType) -> usize {
1004 self._alloc_cursor_id(None, cursor_type)
1005 }
1006
1007 fn _alloc_cursor_id(&mut self, key: Option<CursorKey>, cursor_type: CursorType) -> usize {
1008 let cursor = self.next_free_cursor_id;
1009 self.next_free_cursor_id += 1;
1010 self.cursor_ref.push((key, cursor_type));
1011 turso_assert_eq!(self.cursor_ref.len(), self.next_free_cursor_id);
1012 cursor
1013 }
1014
1015 pub fn add_pragma_result_column(&mut self, col_name: String) {
1016 let expr = ast::Expr::Id(ast::Name::empty());
1019 self.result_columns.push(ResultSetColumn {
1020 expr,
1021 alias: Some(col_name),
1022 implicit_column_name: None,
1023 contains_aggregates: false,
1024 });
1025 }
1026
1027 #[instrument(skip(self), level = Level::DEBUG)]
1028 pub fn emit_insn(&mut self, insn: Insn) {
1029 tracing::trace!("");
1031 self.flags
1032 .set_readonly(self.flags.readonly() & insn.is_readonly());
1033 if matches!(insn, Insn::Function { .. }) {
1035 self.emitted_function_call = true;
1036 }
1037 self.insns.push((insn, self.insns.len()));
1038 }
1039
1040 #[instrument(skip(self), level = Level::DEBUG)]
1044 pub fn emit_no_constant_insn(&mut self, insn: Insn) {
1045 if !self.constant_span_is_open() {
1046 self.constant_span_end_all();
1047 }
1048 self.emit_insn(insn);
1049 }
1050
1051 pub fn close_cursors(&mut self, cursors: &[CursorID]) {
1052 for cursor in cursors {
1053 self.emit_insn(Insn::Close { cursor_id: *cursor });
1054 }
1055 }
1056
1057 pub fn emit_string8(&mut self, value: String, dest: usize) {
1058 self.emit_insn(Insn::String8 { value, dest });
1059 }
1060
1061 pub fn emit_string8_new_reg(&mut self, value: String) -> usize {
1062 let dest = self.alloc_register();
1063 self.emit_insn(Insn::String8 { value, dest });
1064 dest
1065 }
1066
1067 pub fn emit_int(&mut self, value: i64, dest: usize) {
1068 self.emit_insn(Insn::Integer { value, dest });
1069 }
1070
1071 pub fn emit_bool(&mut self, value: bool, dest: usize) {
1072 self.emit_insn(Insn::Integer {
1073 value: if value { 1 } else { 0 },
1074 dest,
1075 });
1076 }
1077
1078 pub fn emit_null(&mut self, dest: usize, dest_end: Option<usize>) {
1079 self.emit_insn(Insn::Null { dest, dest_end });
1080 }
1081
1082 pub fn emit_result_row(&mut self, start_reg: usize, count: usize) {
1083 self.emit_insn(Insn::ResultRow { start_reg, count });
1084 }
1085
1086 fn emit_halt(&mut self, rollback: bool) {
1087 self.emit_insn(Insn::Halt {
1088 err_code: 0,
1089 description: if rollback {
1090 "rollback".to_string()
1091 } else {
1092 String::new()
1093 },
1094 on_error: None,
1095 description_reg: None,
1096 });
1097 }
1098
1099 #[allow(dead_code)]
1103 pub fn emit_halt_err(&mut self, err_code: usize, description: String) {
1104 self.emit_insn(Insn::Halt {
1105 err_code,
1106 description,
1107 on_error: None,
1108 description_reg: None,
1109 });
1110 }
1111
1112 pub fn add_comment(&mut self, insn_index: BranchOffset, comment: &'static str) {
1113 if let QueryMode::Explain | QueryMode::ExplainQueryPlan = self.query_mode {
1114 self.comments.push((insn_index.as_offset_int(), comment));
1115 }
1116 }
1117
1118 pub const fn get_query_mode(&self) -> QueryMode {
1119 self.query_mode
1120 }
1121
1122 pub fn emit_explain(&mut self, push: bool, detail: String) {
1125 if let QueryMode::ExplainQueryPlan = self.query_mode {
1126 self.emit_insn(Insn::Explain {
1127 p1: self.insns.len(),
1128 p2: self.current_parent_explain_idx,
1129 detail,
1130 });
1131 if push {
1132 self.current_parent_explain_idx = Some(self.insns.len() - 1);
1133 }
1134 }
1135 }
1136
1137 pub fn pop_current_parent_explain(&mut self) {
1138 if let QueryMode::ExplainQueryPlan = self.query_mode {
1139 if let Some(current) = self.current_parent_explain_idx {
1140 let (Insn::Explain { p2, .. }, _) = &self.insns[current] else {
1141 unreachable!("current_parent_explain_idx must point to an Explain insn");
1142 };
1143 self.current_parent_explain_idx = *p2;
1144 }
1145 } else {
1146 turso_debug_assert!(self.current_parent_explain_idx.is_none());
1147 }
1148 }
1149
1150 pub fn mark_last_insn_constant(&mut self) {
1151 if self.constant_span_is_open() {
1152 return;
1154 }
1155
1156 let prev = self.insns.len().saturating_sub(1);
1157 self.constant_spans.push((prev, prev));
1158 }
1159
1160 fn emit_constant_insns(&mut self) {
1161 let mut non_constant = Vec::with_capacity(self.insns.len());
1169 let mut constant = Vec::new();
1170 let mut span_idx = 0;
1171
1172 for item in self.insns.drain(..) {
1173 let idx = item.1;
1174
1175 while span_idx < self.constant_spans.len() && self.constant_spans[span_idx].1 < idx {
1177 span_idx += 1;
1178 }
1179
1180 let is_constant =
1182 span_idx < self.constant_spans.len() && self.constant_spans[span_idx].0 <= idx;
1183
1184 if is_constant {
1185 constant.push(item);
1186 } else {
1187 non_constant.push(item);
1188 }
1189 }
1190
1191 self.insns = non_constant;
1192 self.insns.extend(constant);
1193
1194 let mut old_to_new = vec![0usize; self.insns.len()];
1196 for (new_pos, (_, old_idx)) in self.insns.iter().enumerate() {
1197 old_to_new[*old_idx] = new_pos;
1198 }
1199
1200 for resolved_offset in self.label_to_resolved_offset.iter_mut() {
1201 if let Some(old_offset) = resolved_offset {
1202 *resolved_offset = Some(old_to_new[*old_offset as usize] as u32);
1203 }
1204 }
1205
1206 for (offset, _) in self.comments.iter_mut() {
1207 *offset = old_to_new[*offset as usize] as u32;
1208 }
1209
1210 if let QueryMode::ExplainQueryPlan = self.query_mode {
1211 self.current_parent_explain_idx =
1212 self.current_parent_explain_idx.map(|old| old_to_new[old]);
1213
1214 for i in 0..self.insns.len() {
1215 let (Insn::Explain { p2, .. }, _) = &self.insns[i] else {
1216 continue;
1217 };
1218
1219 let new_p2 = p2.map(|old| old_to_new[old]);
1220
1221 let (Insn::Explain { p1, p2, .. }, _) = &mut self.insns[i] else {
1222 unreachable!();
1223 };
1224
1225 *p1 = i;
1226 *p2 = new_p2;
1227 }
1228 }
1229 }
1230
1231 pub const fn offset(&self) -> BranchOffset {
1232 BranchOffset::Offset(self.insns.len() as InsnReference)
1233 }
1234
1235 pub fn allocate_label(&mut self) -> BranchOffset {
1236 let label_n = self.label_to_resolved_offset.len();
1237 self.label_to_resolved_offset.push(None);
1238 BranchOffset::Label(label_n as u32)
1239 }
1240
1241 #[inline]
1253 pub fn preassign_label_to_next_insn(&mut self, label: BranchOffset) {
1254 let BranchOffset::Label(label_number) = label else {
1255 unreachable!("preassign_label_to_next_insn requires a Label, got {label:?}");
1256 };
1257 let anchor = self.offset().as_offset_int().saturating_sub(1);
1258 self.label_to_resolved_offset[label_number as usize] = Some(anchor);
1259 }
1260
1261 #[inline]
1273 pub fn link_label_to_other_label(&mut self, dest: BranchOffset, anchor: BranchOffset) {
1274 let BranchOffset::Label(dest_n) = dest else {
1275 unreachable!("link_label_to_other_label dest must be a Label, got {dest:?}");
1276 };
1277 let BranchOffset::Label(anchor_n) = anchor else {
1278 unreachable!("link_label_to_other_label anchor must be a Label, got {anchor:?}");
1279 };
1280 let resolution = self.label_to_resolved_offset[anchor_n as usize]
1281 .expect("anchor label must already be preassigned/resolved");
1282 self.label_to_resolved_offset[dest_n as usize] = Some(resolution);
1283 }
1284
1285 pub fn resolve_labels(&mut self) -> crate::Result<()> {
1291 let resolve = |pc: &mut BranchOffset, insn_name: &str| -> crate::Result<()> {
1292 if let BranchOffset::Label(label) = pc {
1293 let Some(Some(anchor)) = self.label_to_resolved_offset.get(*label as usize) else {
1294 crate::bail_corrupt_error!(
1295 "Reference to undefined or unresolved label in {insn_name}: {label}"
1296 );
1297 };
1298 *pc = BranchOffset::Offset(anchor + 1);
1299 }
1300 Ok(())
1301 };
1302 for (insn, _) in self.insns.iter_mut() {
1303 match insn {
1304 Insn::Init { target_pc } => {
1305 resolve(target_pc, "Init")?;
1306 }
1307 Insn::Eq {
1308 lhs: _lhs,
1309 rhs: _rhs,
1310 target_pc,
1311 ..
1312 } => {
1313 resolve(target_pc, "Eq")?;
1314 }
1315 Insn::Ne {
1316 lhs: _lhs,
1317 rhs: _rhs,
1318 target_pc,
1319 ..
1320 } => {
1321 resolve(target_pc, "Ne")?;
1322 }
1323 Insn::Lt {
1324 lhs: _lhs,
1325 rhs: _rhs,
1326 target_pc,
1327 ..
1328 } => {
1329 resolve(target_pc, "Lt")?;
1330 }
1331 Insn::Le {
1332 lhs: _lhs,
1333 rhs: _rhs,
1334 target_pc,
1335 ..
1336 } => {
1337 resolve(target_pc, "Le")?;
1338 }
1339 Insn::Gt {
1340 lhs: _lhs,
1341 rhs: _rhs,
1342 target_pc,
1343 ..
1344 } => {
1345 resolve(target_pc, "Gt")?;
1346 }
1347 Insn::Ge {
1348 lhs: _lhs,
1349 rhs: _rhs,
1350 target_pc,
1351 ..
1352 } => {
1353 resolve(target_pc, "Ge")?;
1354 }
1355 Insn::If {
1356 reg: _reg,
1357 target_pc,
1358 jump_if_null: _,
1359 } => {
1360 resolve(target_pc, "If")?;
1361 }
1362 Insn::IfNot {
1363 reg: _reg,
1364 target_pc,
1365 jump_if_null: _,
1366 } => {
1367 resolve(target_pc, "IfNot")?;
1368 }
1369 Insn::Rewind { pc_if_empty, .. } => {
1370 resolve(pc_if_empty, "Rewind")?;
1371 }
1372 Insn::Last { pc_if_empty, .. } => {
1373 resolve(pc_if_empty, "Last")?;
1374 }
1375 Insn::Goto { target_pc } => {
1376 resolve(target_pc, "Goto")?;
1377 }
1378 Insn::DecrJumpZero {
1379 reg: _reg,
1380 target_pc,
1381 } => {
1382 resolve(target_pc, "DecrJumpZero")?;
1383 }
1384 Insn::SorterNext {
1385 cursor_id: _cursor_id,
1386 pc_if_next,
1387 } => {
1388 resolve(pc_if_next, "SorterNext")?;
1389 }
1390 Insn::SorterSort { pc_if_empty, .. } => {
1391 resolve(pc_if_empty, "SorterSort")?;
1392 }
1393 Insn::SorterCompare {
1394 pc_when_nonequal: target_pc,
1395 ..
1396 } => {
1397 resolve(target_pc, "SorterCompare")?;
1398 }
1399 Insn::NotNull {
1400 reg: _reg,
1401 target_pc,
1402 } => {
1403 resolve(target_pc, "NotNull")?;
1404 }
1405 Insn::ColumnHasField { target_pc, .. } => {
1406 resolve(target_pc, "ColumnHasField")?;
1407 }
1408 Insn::IfPos { target_pc, .. } => {
1409 resolve(target_pc, "IfPos")?;
1410 }
1411 Insn::Next { pc_if_next, .. } => {
1412 resolve(pc_if_next, "Next")?;
1413 }
1414 Insn::Once {
1415 target_pc_when_reentered,
1416 ..
1417 } => {
1418 resolve(target_pc_when_reentered, "Once")?;
1419 }
1420 Insn::Prev { pc_if_prev, .. } => {
1421 resolve(pc_if_prev, "Prev")?;
1422 }
1423 Insn::InitCoroutine {
1424 yield_reg: _,
1425 jump_on_definition,
1426 start_offset,
1427 } => {
1428 resolve(jump_on_definition, "InitCoroutine")?;
1429 resolve(start_offset, "InitCoroutine")?;
1430 }
1431 Insn::NotExists {
1432 cursor: _,
1433 rowid_reg: _,
1434 target_pc,
1435 } => {
1436 resolve(target_pc, "NotExists")?;
1437 }
1438 Insn::MustBeInt {
1439 target_pc: Some(target_pc),
1440 ..
1441 } => {
1442 resolve(target_pc, "MustBeInt")?;
1443 }
1444 Insn::Yield {
1445 yield_reg: _,
1446 end_offset,
1447 subtype_clear_start_reg: _,
1448 subtype_clear_count: _,
1449 } => {
1450 resolve(end_offset, "Yield")?;
1451 }
1452 Insn::SeekRowid { target_pc, .. } => {
1453 resolve(target_pc, "SeekRowid")?;
1454 }
1455 Insn::Gosub { target_pc, .. } => {
1456 resolve(target_pc, "Gosub")?;
1457 }
1458 Insn::Jump {
1459 target_pc_eq,
1460 target_pc_lt,
1461 target_pc_gt,
1462 } => {
1463 resolve(target_pc_eq, "Jump")?;
1464 resolve(target_pc_lt, "Jump")?;
1465 resolve(target_pc_gt, "Jump")?;
1466 }
1467 Insn::SeekGE { target_pc, .. } => resolve(target_pc, "SeekGE")?,
1468 Insn::SeekGT { target_pc, .. } => resolve(target_pc, "SeekGT")?,
1469 Insn::SeekLE { target_pc, .. } => resolve(target_pc, "SeekLE")?,
1470 Insn::SeekLT { target_pc, .. } => resolve(target_pc, "SeekLT")?,
1471 Insn::IdxGE { target_pc, .. } => resolve(target_pc, "IdxGE")?,
1472 Insn::IdxLE { target_pc, .. } => resolve(target_pc, "IdxLE")?,
1473 Insn::IdxGT { target_pc, .. } => resolve(target_pc, "IdxGT")?,
1474 Insn::IdxLT { target_pc, .. } => resolve(target_pc, "IdxLT")?,
1475 Insn::IndexMethodQuery { pc_if_empty, .. } => {
1476 resolve(pc_if_empty, "IndexMethodQuery")?;
1477 }
1478 Insn::IsNull { reg: _, target_pc } => resolve(target_pc, "IsNull")?,
1479 Insn::VNext { pc_if_next, .. } => resolve(pc_if_next, "VNext")?,
1480 Insn::VFilter { pc_if_empty, .. } => resolve(pc_if_empty, "VFilter")?,
1481 Insn::RowSetRead { pc_if_empty, .. } => resolve(pc_if_empty, "RowSetRead")?,
1482 Insn::RowSetTest { pc_if_found, .. } => resolve(pc_if_found, "RowSetTest")?,
1483 Insn::NoConflict { target_pc, .. } => resolve(target_pc, "NoConflict")?,
1484 Insn::Found { target_pc, .. } => resolve(target_pc, "Found")?,
1485 Insn::NotFound { target_pc, .. } => resolve(target_pc, "NotFound")?,
1486 Insn::FkIfZero { target_pc, .. } => resolve(target_pc, "FkIfZero")?,
1487 Insn::Filter { target_pc, .. } => resolve(target_pc, "Filter")?,
1488 Insn::HashProbe { target_pc, .. } => resolve(target_pc, "HashProbe")?,
1489 Insn::HashNext { target_pc, .. } => resolve(target_pc, "HashNext")?,
1490 Insn::HashDistinct { data } => resolve(&mut data.target_pc, "HashDistinct")?,
1491 Insn::HashScanUnmatched { target_pc, .. } => {
1492 resolve(target_pc, "HashScanUnmatched")?
1493 }
1494 Insn::HashNextUnmatched { target_pc, .. } => {
1495 resolve(target_pc, "HashNextUnmatched")?
1496 }
1497 Insn::HashGraceInit { target_pc, .. } => resolve(target_pc, "HashGraceInit")?,
1498 Insn::HashGraceLoadPartition { target_pc, .. } => {
1499 resolve(target_pc, "HashGraceLoadPartition")?
1500 }
1501 Insn::HashGraceNextProbe { target_pc, .. } => {
1502 resolve(target_pc, "HashGraceNextProbe")?
1503 }
1504 Insn::HashGraceAdvancePartition { target_pc, .. } => {
1505 resolve(target_pc, "HashGraceAdvancePartition")?
1506 }
1507 Insn::Program {
1508 ignore_jump_target, ..
1509 } => resolve(ignore_jump_target, "Program")?,
1510 _ => {}
1511 }
1512 }
1513 self.label_to_resolved_offset.clear();
1514 Ok(())
1515 }
1516
1517 pub fn set_cursor_override(&mut self, table_ref_id: TableInternalId, cursor_id: CursorID) {
1520 self.cursor_overrides.insert(table_ref_id.into(), cursor_id);
1521 }
1522
1523 pub fn clear_cursor_override(&mut self, table_ref_id: TableInternalId) {
1525 self.cursor_overrides.remove(&table_ref_id.into());
1526 }
1527
1528 pub fn clear_all_cursor_overrides(&mut self) {
1530 self.cursor_overrides.clear();
1531 }
1532
1533 pub fn has_cursor_override(&self, table_ref_id: TableInternalId) -> bool {
1535 self.cursor_overrides.contains_key(&table_ref_id.into())
1536 }
1537
1538 pub fn resolve_cursor_id_safe(&self, key: &CursorKey) -> Option<CursorID> {
1540 if key.index.is_none() && !key.is_build {
1544 let table_id: usize = key.table_reference_id.into();
1545 if let Some(&cursor_id) = self.cursor_overrides.get(&table_id) {
1546 return Some(cursor_id);
1547 }
1548 }
1549 self.cursor_ref
1550 .iter()
1551 .position(|(k, _)| k.as_ref().is_some_and(|k| k.equals(key)))
1552 }
1553
1554 pub fn resolve_cursor_id(&self, key: &CursorKey) -> CursorID {
1555 self.resolve_cursor_id_safe(key)
1556 .unwrap_or_else(|| panic!("Cursor not found: {key:?}"))
1557 }
1558
1559 pub fn resolve_any_index_cursor_id_for_table(&self, table_ref_id: TableInternalId) -> CursorID {
1566 self.resolve_any_index_cursor_id_for_table_safe(table_ref_id)
1567 .unwrap_or_else(|| panic!("No index cursor found for table {table_ref_id}"))
1568 }
1569
1570 pub fn resolve_any_index_cursor_id_for_table_safe(
1571 &self,
1572 table_ref_id: TableInternalId,
1573 ) -> Option<CursorID> {
1574 self.cursor_ref.iter().position(|(k, _)| {
1575 k.as_ref()
1576 .is_some_and(|k| k.table_reference_id == table_ref_id && k.index.is_some())
1577 })
1578 }
1579
1580 pub fn resolve_index_for_cursor_id(&self, cursor_id: CursorID) -> Arc<Index> {
1582 let cursor_ref = &self
1583 .cursor_ref
1584 .get(cursor_id)
1585 .unwrap_or_else(|| panic!("Cursor not found: {cursor_id}"))
1586 .1;
1587 let CursorType::BTreeIndex(index) = cursor_ref else {
1588 panic!("Cursor is not an index: {cursor_id}");
1589 };
1590 index.clone()
1591 }
1592
1593 pub fn get_cursor_type(&self, cursor_id: CursorID) -> Option<&CursorType> {
1595 self.cursor_ref
1596 .get(cursor_id)
1597 .map(|(_, cursor_type)| cursor_type)
1598 }
1599
1600 pub const fn set_collation(&mut self, c: Option<(CollationSeq, bool)>) {
1601 self.collation = c
1602 }
1603
1604 pub const fn curr_collation_ctx(&self) -> Option<(CollationSeq, bool)> {
1605 self.collation
1606 }
1607
1608 pub fn curr_collation(&self) -> Option<CollationSeq> {
1609 self.collation.map(|c| c.0)
1610 }
1611
1612 pub const fn reset_collation(&mut self) {
1613 self.collation = None;
1614 }
1615
1616 #[inline]
1617 pub fn nested<T>(&mut self, body: impl FnOnce(&mut Self) -> T) -> T {
1618 self.incr_nesting();
1619 let res = body(self);
1620 self.decr_nesting();
1621 res
1622 }
1623
1624 #[inline]
1625 const fn incr_nesting(&mut self) {
1626 self.nested_level += 1;
1627 }
1628
1629 #[inline]
1630 const fn decr_nesting(&mut self) {
1631 self.nested_level -= 1;
1632 }
1633
1634 #[inline]
1636 pub const fn is_nested(&self) -> bool {
1637 self.nested_level > 0
1638 }
1639
1640 pub fn prologue(&mut self) {
1642 if self.flags.is_subprogram() {
1643 self.init_label = self.allocate_label();
1645 self.emit_insn(Insn::Init {
1646 target_pc: self.init_label,
1647 });
1648 self.preassign_label_to_next_insn(self.init_label);
1649 self.start_offset = self.offset();
1650 return;
1651 }
1652 if self.nested_level == 0 {
1653 self.init_label = self.allocate_label();
1654
1655 self.emit_insn(Insn::Init {
1656 target_pc: self.init_label,
1657 });
1658
1659 self.start_offset = self.offset();
1660 }
1661 }
1662
1663 pub fn begin_write_operation(&mut self) -> Result<(), alloc::TryReserveError> {
1665 self.txn_mode = TransactionMode::Write;
1666 self.write_databases.set(crate::MAIN_DB_ID)
1667 }
1668
1669 pub fn begin_write_on_database(
1671 &mut self,
1672 database_id: usize,
1673 schema_cookie: u32,
1674 ) -> Result<(), alloc::TryReserveError> {
1675 self.txn_mode = TransactionMode::Write;
1676 self.write_databases.set(database_id)?;
1677 self.write_database_cookies
1678 .insert(database_id, schema_cookie);
1679 Ok(())
1680 }
1681
1682 pub fn begin_read_operation(&mut self) -> Result<(), alloc::TryReserveError> {
1683 if matches!(self.txn_mode, TransactionMode::None) {
1685 self.txn_mode = TransactionMode::Read;
1686 }
1687 self.read_databases.set(crate::MAIN_DB_ID)
1688 }
1689
1690 pub fn begin_read_on_database(
1694 &mut self,
1695 database_id: usize,
1696 schema_cookie: u32,
1697 ) -> Result<(), alloc::TryReserveError> {
1698 if matches!(self.txn_mode, TransactionMode::None) {
1699 self.txn_mode = TransactionMode::Read;
1700 }
1701 self.read_databases.set(database_id)?;
1702 self.read_database_cookies
1703 .insert(database_id, schema_cookie);
1704 Ok(())
1705 }
1706
1707 pub const fn begin_concurrent_operation(&mut self) {
1708 self.txn_mode = TransactionMode::Concurrent;
1709 }
1710
1711 pub const fn rollback(&mut self) {
1713 self.flags.set_rollback(true);
1714 }
1715
1716 pub fn epilogue(&mut self, schema: &Schema) {
1720 if self.flags.is_subprogram() {
1721 let description = if self.trigger.is_some() {
1723 "trigger"
1724 } else {
1725 "fk action"
1726 };
1727 self.emit_insn(Insn::Halt {
1728 err_code: 0,
1729 description: description.to_string(),
1730 on_error: None,
1731 description_reg: None,
1732 });
1733 return;
1734 }
1735 if self.nested_level == 0 {
1736 self.emit_halt(self.flags.rollback());
1738 self.preassign_label_to_next_insn(self.init_label);
1739
1740 if !matches!(self.txn_mode, TransactionMode::None) {
1741 let write_dbs = self.write_databases.clone();
1742 for db_id in &write_dbs {
1743 let schema_cookie = if db_id == crate::MAIN_DB_ID {
1744 schema.schema_version
1745 } else {
1746 self.write_database_cookies
1747 .get(&db_id)
1748 .copied()
1749 .unwrap_or(0)
1750 };
1751 self.emit_insn(Insn::Transaction {
1752 db: db_id,
1753 tx_mode: self.txn_mode,
1754 schema_cookie,
1755 });
1756 }
1757 let read_dbs = self.read_databases.clone();
1760 for db_id in &read_dbs {
1761 if !write_dbs.get(db_id) {
1762 let schema_cookie = if db_id == crate::MAIN_DB_ID {
1763 schema.schema_version
1764 } else {
1765 self.read_database_cookies.get(&db_id).copied().unwrap_or(0)
1766 };
1767 self.emit_insn(Insn::Transaction {
1768 db: db_id,
1769 tx_mode: TransactionMode::Read,
1770 schema_cookie,
1771 });
1772 }
1773 }
1774 }
1775
1776 if !self.constant_spans.is_empty() {
1777 self.emit_constant_insns();
1778 }
1779 self.emit_insn(Insn::Goto {
1780 target_pc: self.start_offset,
1781 });
1782 }
1783 }
1784
1785 pub fn is_table_open(&self, table: &Table) -> bool {
1787 self.table_references.contains_table(table)
1788 }
1789
1790 pub fn cursor_is_btree(&self, cursor_id: CursorID) -> bool {
1792 matches!(self.cursor_ref[cursor_id].1, CursorType::BTreeTable(_))
1793 }
1794
1795 pub fn btree_table_from_cursor(&self, cursor_id: CursorID) -> Option<&Arc<BTreeTable>> {
1797 match &self.cursor_ref[cursor_id].1 {
1798 CursorType::BTreeTable(t) => Some(t),
1799 _ => None,
1800 }
1801 }
1802
1803 #[inline]
1804 pub fn cursor_loop(&mut self, cursor_id: CursorID, f: impl Fn(&mut ProgramBuilder, usize)) {
1805 let loop_start = self.allocate_label();
1806 let loop_end = self.allocate_label();
1807
1808 self.emit_insn(Insn::Rewind {
1809 cursor_id,
1810 pc_if_empty: loop_end,
1811 });
1812 self.preassign_label_to_next_insn(loop_start);
1813
1814 let rowid = self.alloc_register();
1815
1816 self.emit_insn(Insn::RowId {
1817 cursor_id,
1818 dest: rowid,
1819 });
1820
1821 self.emit_insn(Insn::IsNull {
1822 reg: rowid,
1823 target_pc: loop_end,
1824 });
1825
1826 f(self, rowid);
1827
1828 self.emit_insn(Insn::Next {
1829 cursor_id,
1830 pc_if_next: loop_start,
1831 });
1832 self.preassign_label_to_next_insn(loop_end);
1833 }
1834
1835 pub fn emit_column_or_rowid(&mut self, cursor_id: CursorID, column: usize, out: usize) {
1836 let (_, cursor_type) = self.cursor_ref.get(cursor_id).expect("cursor_id is valid");
1837 if let CursorType::BTreeTable(btree) = cursor_type {
1838 let column_def = btree
1839 .columns()
1840 .get(column)
1841 .expect("column index out of bounds");
1842 if column_def.is_rowid_alias() {
1843 self.flags.set_suppress_column_default(false);
1846 self.emit_insn(Insn::RowId {
1847 cursor_id,
1848 dest: out,
1849 });
1850 } else {
1851 self.emit_column(cursor_id, column, out);
1852 }
1853 } else {
1854 self.emit_column(cursor_id, column, out);
1855 }
1856 }
1857
1858 pub fn emit_column_has_field(
1862 &mut self,
1863 cursor_id: CursorID,
1864 column: usize,
1865 target_pc: BranchOffset,
1866 ) {
1867 let (_, cursor_type) = self.cursor_ref.get(cursor_id).expect("cursor_id is valid");
1868 let physical_column = match cursor_type {
1869 CursorType::BTreeTable(btree) => btree.logical_to_physical_column(column),
1870 _ => column,
1871 };
1872 self.emit_insn(Insn::ColumnHasField {
1873 cursor_id,
1874 column: physical_column,
1875 target_pc,
1876 });
1877 }
1878
1879 pub fn emit_column_affinity(&mut self, register: usize, affinity: Affinity) {
1881 self.emit_insn(Insn::Affinity {
1882 start_reg: register,
1883 count: NonZeroUsize::MIN,
1884 affinities: affinity.aff_mask().to_string(),
1885 });
1886 }
1887
1888 fn emit_column(&mut self, cursor_id: CursorID, column: usize, out: usize) {
1889 let (_, cursor_type) = self.cursor_ref.get(cursor_id).expect("cursor_id is valid");
1890
1891 if let CursorType::BTreeTable(btree) = cursor_type {
1892 let column_def = btree
1893 .columns()
1894 .get(column)
1895 .expect("column index out of bounds");
1896 turso_assert!(
1897 !column_def.is_virtual_generated(),
1898 "emit_column called with virtual generated column index",
1899 {"column_index": column}
1900 );
1901 }
1902
1903 let physical_column = match cursor_type {
1904 CursorType::BTreeTable(btree) => btree.logical_to_physical_column(column),
1905 _ => column,
1906 };
1907
1908 let default = 'value: {
1909 let default = match cursor_type {
1910 CursorType::BTreeTable(btree) => &btree.columns()[column].default,
1911 CursorType::BTreeIndex(index) => &index.columns[column].default,
1912 CursorType::MaterializedView(btree, _) => &btree.columns()[column].default,
1913 _ => break 'value None,
1914 };
1915
1916 let Some(ref default_expr) = default else {
1917 break 'value None;
1918 };
1919
1920 let mut value = match crate::translate::alter::eval_constant_default_value(default_expr)
1927 {
1928 Ok(v) => v,
1929 Err(_) => break 'value None,
1930 };
1931
1932 let affinity = match cursor_type {
1937 CursorType::BTreeTable(btree) => btree.columns()[column].affinity(),
1938 CursorType::MaterializedView(btree, _) => btree.columns()[column].affinity(),
1939 _ => Affinity::Blob,
1940 };
1941 if let Some(converted) = affinity.convert(&value) {
1942 value = match converted {
1943 either::Either::Left(val_ref) => val_ref.to_owned(),
1944 either::Either::Right(val) => val,
1945 };
1946 }
1947
1948 Some(value)
1949 };
1950
1951 let default = if self.flags.suppress_column_default() {
1952 self.flags.set_suppress_column_default(false);
1953 None
1954 } else {
1955 default
1956 };
1957
1958 self.emit_insn(Insn::Column {
1959 cursor_id,
1960 column: physical_column,
1961 dest: out,
1962 default,
1963 });
1964 }
1965
1966 pub fn build_prepared_program(
1967 mut self,
1968 prepare_context: PrepareContext,
1969 change_cnt_on: bool,
1970 sql: &str,
1971 ) -> crate::Result<PreparedProgram> {
1972 self.resolve_labels()?;
1973
1974 self.parameters.list.dedup();
1975
1976 let needs_stmt_subtransactions = matches!(self.txn_mode, TransactionMode::Write)
1982 && self.flags.is_multi_write()
1983 && self.may_abort();
1984
1985 let prepared = PreparedProgram {
1986 max_registers: self.next_free_register,
1987 insns: self.insns,
1988 cursor_ref: self.cursor_ref,
1989 comments: self.comments,
1990 parameters: self.parameters,
1991 change_cnt_on,
1992 readonly: self.flags.readonly(),
1993 result_columns: self.result_columns,
1994 table_references: self.table_references,
1995 sql: sql.to_string(),
1996 needs_stmt_subtransactions: crate::Arc::new(crate::AtomicBool::new(
1997 needs_stmt_subtransactions,
1998 )),
1999 trigger: self.trigger.take(),
2000 is_subprogram: self.flags.is_subprogram(),
2001 resolve_type: self.resolve_type,
2002 prepare_context,
2003 write_databases: self.write_databases,
2004 read_databases: self.read_databases,
2005 };
2006 Ok(prepared)
2007 }
2008
2009 #[turso_macros::trace_stack]
2010 pub fn build(
2011 self,
2012 connection: Arc<Connection>,
2013 change_cnt_on: bool,
2014 sql: &str,
2015 ) -> crate::Result<Program> {
2016 let prepare_context = PrepareContext::from_connection(&connection);
2017 let prepared = self.build_prepared_program(prepare_context, change_cnt_on, sql)?;
2018 Ok(Program::from_prepared(Arc::new(prepared), connection))
2019 }
2020}