// Copyright 2024 Vincent Chan // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. use super::label::{JumpTableRecord, Label, LabelSlot}; use crate::coll::collection_info::CollectionSpecification; use crate::errors::{mk_invalid_query_field}; use crate::index::INDEX_PREFIX; use crate::vm::op::DbOp; use crate::vm::subprogram::SubProgramIndexItem; use crate::vm::SubProgram; use crate::{Error, Result}; use bson::spec::{BinarySubtype, ElementType}; use bson::{Array, Binary, Bson, Document}; use crate::vm::aggregation_codegen_context::{AggregationCodeGenContext, PipelineItem}; use crate::vm::global_variable::{GlobalVariable, GlobalVariableSlot}; use crate::vm::operators::OpRegistry; use crate::vm::update_operators::{IncOperator, MaxOperator, MinOperator, MulOperator, PopOperator, PushOperator, RenameOperator, SetOperator, UnsetOperator, UpdateOperator}; use crate::vm::vm_add_fields::VmFuncAddFields; use crate::vm::vm_count::VmFuncCount; use crate::vm::vm_external_func::VmExternalFunc; use crate::vm::vm_group::VmFuncGroup; use crate::vm::vm_limit::VmFuncLimit; use crate::vm::vm_skip::VmFuncSkip; use crate::vm::vm_sort::VmFuncSort; use crate::vm::vm_unset::VmFuncUnset; const JUMP_TABLE_DEFAULT_SIZE: usize = 8; const PATH_DEFAULT_SIZE: usize = 8; pub(super) struct Codegen { program: Box, jump_table: Vec, skip_annotation: bool, is_write: bool, paths: Vec, op_registry: OpRegistry, } impl Codegen { pub(super) fn new(skip_annotation: bool, is_write: bool) -> Codegen { Codegen { program: Box::new(SubProgram::new()), jump_table: Vec::with_capacity(JUMP_TABLE_DEFAULT_SIZE), skip_annotation, is_write, paths: Vec::with_capacity(PATH_DEFAULT_SIZE), op_registry: OpRegistry, } } fn unify_labels(&mut self) { for record in &self.jump_table { let pos = (record.begin_loc + record.offset) as usize; let slot = &self.program.label_slots[record.label_id as usize]; let target = slot.position(); let bytes: [u8; 4] = target.to_le_bytes(); self.program.instructions[pos..pos + 4].copy_from_slice(&bytes); } } pub(super) fn take(mut self) -> SubProgram { self.unify_labels(); *self.program } #[inline] #[allow(dead_code)] pub(super) fn new_global_variable(&mut self, init_value: Bson) -> Result { self.new_global_variable_impl(init_value, None) } #[inline] #[allow(dead_code)] pub(super) fn new_global_variable_with_name(&mut self, name: String, init_value: Bson) -> Result { self.new_global_variable_impl(init_value, Some(name.into_boxed_str())) } fn new_global_variable_impl(&mut self, init_value: Bson, name: Option>) -> Result { let id = self.program.global_variables.len() as u32; self.program.global_variables.push(GlobalVariableSlot { pos: 0, init_value, name, }); Ok(GlobalVariable::new(id)) } pub(super) fn new_label(&mut self) -> Label { let id = self.program.label_slots.len() as u32; self.program.label_slots.push(LabelSlot::Empty); Label::new(id) } pub(super) fn emit_label(&mut self, label: Label) { if !self.program.label_slots[label.u_pos()].is_empty() { unreachable!("this label has been emit"); } let current_loc = self.current_location(); self.emit(DbOp::Label); self.emit_u32(label.pos()); self.program.label_slots[label.u_pos()] = LabelSlot::UnnamedLabel(current_loc); } #[allow(dead_code)] fn emit_load_global(&mut self, global: GlobalVariable) { self.emit(DbOp::LoadGlobal); self.emit_u32(global.pos()); } #[allow(dead_code)] fn emit_store_global(&mut self, global: GlobalVariable) { self.emit(DbOp::StoreGlobal); self.emit_u32(global.pos()); } pub(super) fn emit_label_with_name>>(&mut self, label: Label, name: T) { if !self.program.label_slots[label.u_pos()].is_empty() { unreachable!("this label has been emit"); } let current_loc = self.current_location(); self.emit(DbOp::Label); self.emit_u32(label.pos()); if self.skip_annotation { self.program.label_slots[label.u_pos()] = LabelSlot::UnnamedLabel(current_loc); } else { self.program.label_slots[label.u_pos()] = LabelSlot::LabelWithString(current_loc, name.into()); } } fn emit_query_layout_has_pkey( &mut self, pkey: Bson, query: &Document, result_callback: F, ) -> Result<()> where F: FnOnce(&mut Codegen) -> Result<()>, { let close_label = self.new_label(); let result_label = self.new_label(); let pkey_id = self.push_static(pkey); self.emit_push_value(pkey_id); self.emit_goto(DbOp::FindByPrimaryKey, close_label); self.emit_goto(DbOp::Goto, result_label); self.emit_label(close_label); self.emit(DbOp::Pop); self.emit(DbOp::Close); self.emit(DbOp::Halt); self.emit_label(result_label); for (key, value) in query.iter() { if key == "_id" { continue; } let key_static_id = self.push_static(Bson::String(key.clone())); let value_static_id = self.push_static(value.clone()); self.emit_goto2(DbOp::GetField, key_static_id, close_label); // push a value1 self.emit_push_value(value_static_id); // push a value2 self.emit(DbOp::Equal); // if not equal,go to next self.emit_goto(DbOp::IfFalse, close_label); self.emit(DbOp::Pop); // pop a value2 self.emit(DbOp::Pop); // pop a value1 } result_callback(self)?; self.emit_goto(DbOp::Goto, close_label); Ok(()) } pub(super) fn emit_query_layout( &mut self, col_spec: &CollectionSpecification, query: &Document, result_callback: F, before_close: Option Result<()>>>, is_many: bool, ) -> Result<()> where F: FnOnce(&mut Codegen) -> Result<()>, { let try_pkey_result = self.try_query_by_pkey(col_spec, query, result_callback)?; if try_pkey_result.is_none() { return Ok(()); } let result_callback: F = try_pkey_result.unwrap(); let try_index_result = self.try_query_by_index(col_spec, query, result_callback)?; if try_index_result.is_none() { return Ok(()); } self.emit_open(col_spec._id.clone().into()); let result_callback: F = try_index_result.unwrap(); let compare_fun = self.new_label(); let compare_fun_clean = self.new_label(); let compare_label = self.new_label(); let next_label = self.new_label(); let result_label = self.new_label(); let not_found_label = self.new_label(); let close_label = self.new_label(); self.emit_goto(DbOp::Rewind, close_label); self.emit_goto(DbOp::Goto, compare_label); self.emit_label(next_label); self.emit_goto(DbOp::Next, compare_label); // <==== close cursor self.emit_label_with_name(close_label, "close"); if let Some(before_close) = before_close { before_close(self)?; } self.emit(DbOp::Close); self.emit(DbOp::Halt); // <==== not this item, go to next item self.emit_label_with_name(not_found_label, "not_this_item"); self.emit(DbOp::Pop); // pop the current value; self.emit_goto(DbOp::Goto, next_label); // <==== result position // give out the result, or update the item self.emit_label_with_name(result_label, "result"); result_callback(self)?; if is_many { self.emit_goto(DbOp::Goto, next_label); } else { self.emit_goto(DbOp::Goto, close_label); } // <==== begin to compare the top of the stack // // the top of the stack is the target document // // begin to execute compare logic // save the stack first self.emit_label_with_name(compare_label, "compare"); self.emit(DbOp::Dup); self.emit_goto(DbOp::Call, compare_fun); self.emit_u32(1); self.emit_goto(DbOp::IfFalse, not_found_label); self.emit_goto(DbOp::Goto, result_label); self.emit_label_with_name(compare_fun, "compare_function"); self.emit_standard_query_doc(query, result_label, compare_fun_clean)?; self.emit_label_with_name(compare_fun_clean, "compare_function_clean"); self.emit_ret(0); Ok(()) } fn try_query_by_pkey( &mut self, col_spec: &CollectionSpecification, query: &Document, result_callback: F, ) -> Result> where F: FnOnce(&mut Codegen) -> Result<()>, { if let Some(id_value) = query.get("_id") { if id_value.element_type() != ElementType::EmbeddedDocument { self.emit_open(col_spec._id.clone().into()); self.emit_query_layout_has_pkey(id_value.clone(), query, result_callback)?; return Ok(None); } } Ok(Some(result_callback)) } fn try_query_by_index( &mut self, col_spec: &CollectionSpecification, query: &Document, result_callback: F, ) -> Result> where F: FnOnce(&mut Codegen) -> Result<()>, { if self.is_write { return Ok(Some(result_callback)); } let index_meta = &col_spec.indexes; for (index_name, index_info) in index_meta { let (key, _order) = index_info.keys.iter().next().unwrap(); // the key is ellipse representation, such as "a.b.c" // the query is supposed to be ellipse too, such as // { "a.b.c": 1 } let test_result = query.get(key); if let Some(query_doc) = test_result { if query_doc.element_type() != ElementType::EmbeddedDocument { let mut remain_query = query.clone(); remain_query.remove(key); self.indeed_emit_query_by_index( col_spec._id.as_str(), index_name.as_str(), query_doc, &remain_query, result_callback, )?; return Ok(None); } } } Ok(Some(result_callback)) } fn indeed_emit_query_by_index( &mut self, col_name: &str, index_name: &str, query_value: &Bson, remain_query: &Document, result_callback: F, ) -> Result<()> where F: FnOnce(&mut Codegen) -> Result<()>, { let prefix_bytes = { let b_prefix = Bson::String(INDEX_PREFIX.to_string()); let b_col_name = Bson::String(col_name.to_string()); let b_index_name = &Bson::String(index_name.to_string()); let buf: Vec<&Bson> = vec![&b_prefix, &b_col_name, &b_index_name]; crate::utils::bson::stacked_key(buf)? }; self.emit_open(Bson::Binary(Binary { subtype: BinarySubtype::Generic, bytes: prefix_bytes, })); let close_label = self.new_label(); let result_label = self.new_label(); let next_label = self.new_label(); let value_id = self.push_static(query_value.clone()); self.emit_push_value(value_id); let col_name_id = self.push_static(Bson::String(col_name.to_string())); self.emit_push_value(col_name_id); self.emit_goto(DbOp::FindByIndex, close_label); self.emit_goto(DbOp::Goto, result_label); self.emit_label(next_label); self.emit_goto(DbOp::NextIndexValue, result_label); self.emit_label(close_label); self.emit(DbOp::Pop); // pop the collection name self.emit(DbOp::Pop); // pop the query value self.emit(DbOp::Close); self.emit(DbOp::Halt); self.emit_label(result_label); for (key, value) in remain_query.iter() { let key_static_id = self.push_static(Bson::String(key.clone())); let value_static_id = self.push_static(value.clone()); self.emit_goto2(DbOp::GetField, key_static_id, close_label); // push a value1 self.emit_push_value(value_static_id); // push a value2 self.emit(DbOp::Equal); // if not equal,go to next self.emit_goto(DbOp::IfFalse, close_label); self.emit(DbOp::Pop); // pop a value2 self.emit(DbOp::Pop); // pop a value1 } result_callback(self)?; self.emit_goto(DbOp::Goto, next_label); Ok(()) } fn emit_standard_query_doc( &mut self, query_doc: &Document, result_label: Label, not_found_label: Label, ) -> Result<()> { if query_doc.is_empty() { self.emit(DbOp::StoreR0_2); self.emit_u8(1); return Ok(()) } for (key, value) in query_doc.iter() { crate::path_hint!(self, key.clone(), { self.emit_query_tuple( key, value, result_label, not_found_label, )?; }); } Ok(()) } fn gen_path(&self) -> String { let mut result = String::with_capacity(32); for item in &self.paths { result.push('/'); result.push_str(item.as_ref()); } result } #[inline] fn last_key(&self) -> &str { self.paths.last().unwrap().as_str() } fn emit_logic_and( &mut self, arr: &Array, result_label: Label, not_found_label: Label, ) -> Result<()> { for (index, item_doc_value) in arr.iter().enumerate() { let path_msg = format!("[{}]", index); crate::path_hint!(self, path_msg, { let item_doc = crate::try_unwrap_document!("$and", item_doc_value); self.emit_standard_query_doc( item_doc, result_label, not_found_label, )?; }); } Ok(()) } fn emit_logic_or( &mut self, arr: &Array, ret_label: Label, ) -> Result<()> { let cmp_label = self.new_label(); self.emit_goto(DbOp::Goto, cmp_label); let mut functions = Vec::