diff --git a/src/codec/postings/mod.rs b/src/codec/postings/mod.rs index d439b64e6..bd1c00d66 100644 --- a/src/codec/postings/mod.rs +++ b/src/codec/postings/mod.rs @@ -93,6 +93,19 @@ pub trait PostingsSerializer { /// blockwand is disabled), the term_doc_freq passed will be set 0. fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool); + /// Codec-specific per-term payload. + /// + /// It is supplied right after `new_term` and before any `write_doc`, so the + /// codec can let it influence how the postings list is encoded. + /// + /// Hidden contract: `new_term` MUST reset any per-term payload state to its + /// default. This method is only called for terms that actually have a + /// payload registered, so a codec cannot rely on it being called for every + /// term. + /// + /// The default implementation ignores the payload. + fn set_term_payload(&mut self, _payload: &dyn std::any::Any) {} + /// Records a new document id for the current term. /// The serializer may ignore it. fn write_doc(&mut self, doc_id: DocId, term_freq: u32); diff --git a/src/docset.rs b/src/docset.rs index 187a3d376..7f5a08ff2 100644 --- a/src/docset.rs +++ b/src/docset.rs @@ -1,8 +1,6 @@ use std::ops::{Deref as _, DerefMut as _}; -use common::BitSet; - -use common::TinySet; +use common::{BitSet, TinySet}; use crate::fastfield::AliveBitSet; use crate::DocId; diff --git a/src/indexer/segment_writer.rs b/src/indexer/segment_writer.rs index e6846180e..fb107444c 100644 --- a/src/indexer/segment_writer.rs +++ b/src/indexer/segment_writer.rs @@ -1,3 +1,5 @@ +use std::any::Any; + use columnar::MonotonicallyMappableToU64; use common::JsonPathWriter; use itertools::Itertools; @@ -16,7 +18,7 @@ use crate::postings::{ PerFieldPostingsWriter, PostingsWriter, PostingsWriterEnum, }; use crate::schema::document::{Document, Value}; -use crate::schema::{FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED}; +use crate::schema::{Field, FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED}; use crate::tokenizer::{FacetTokenizer, PreTokenizedStream, TextAnalyzer, Tokenizer}; use crate::{DocId, Opstamp, TantivyError}; @@ -49,7 +51,7 @@ fn compute_initial_table_size(per_thread_memory_budget: usize) -> crate::Result< pub struct SegmentWriter { pub(crate) max_doc: DocId, pub(crate) ctx: IndexingContext, - pub(crate) per_field_postings_writers: PerFieldPostingsWriter, + pub per_field_postings_writers: PerFieldPostingsWriter, pub(crate) segment_serializer: SegmentSerializer, pub(crate) fast_field_writers: FastFieldsWriter, pub(crate) fieldnorms_writer: FieldNormsWriter, @@ -148,6 +150,67 @@ impl SegmentWriter { + self.segment_serializer.mem_usage() } + /// Attaches or updates a codec-specific payload on a term of a regular + /// (non-JSON) field. + /// + /// `value_bytes` is the serialized term value, i.e. exactly what would be + /// appended after the field id (the raw text bytes for a str field, or the + /// big-endian bytes for a numeric field). + /// + /// If the term does not exist yet, it is inserted with an empty recorder so + /// that it still gets serialized even though it belongs to no document. + /// `updater` receives the previously registered payload (`None` if absent) + /// and returns the payload to store. The payload is handed to the codec's + /// postings serializer (via `set_term_payload`) at the beginning of the + /// term during serialization. + pub(crate) fn update_term_payload( + &mut self, + field: Field, + value_bytes: &[u8], + updater: impl FnOnce(Option>) -> Box, + ) { + let mut term = IndexingTerm::with_capacity(value_bytes.len()); + term.set_field(field); + term.append_bytes(value_bytes); + self.update_term_payload_for_serialized_term(field, term.serialized_term(), updater); + } + + /// Same as [`Self::update_term_payload`] for a JSON field. + /// + /// `value_bytes` must be the type-tagged value (`[type code][value]`), the + /// representation that follows the path within a JSON term. + pub(crate) fn update_json_term_payload( + &mut self, + field: Field, + json_path: &str, + value_bytes: &[u8], + updater: impl FnOnce(Option>) -> Box, + ) { + let unordered_id = self + .ctx + .path_to_unordered_id + .get_or_allocate_unordered_id(json_path); + // JSON term key layout: `[field:4][unordered_path_id:4][type code][value]`. + let mut serialized_term = Vec::with_capacity(8 + value_bytes.len()); + serialized_term.extend_from_slice(&field.field_id().to_be_bytes()); + serialized_term.extend_from_slice(&unordered_id.to_be_bytes()); + serialized_term.extend_from_slice(value_bytes); + self.update_term_payload_for_serialized_term(field, &serialized_term, updater); + } + + fn update_term_payload_for_serialized_term( + &mut self, + field: Field, + serialized_term: &[u8], + updater: impl FnOnce(Option>) -> Box, + ) { + let postings_writer = self.per_field_postings_writers.get_for_field(field); + let addr = postings_writer.ensure_term(serialized_term, &mut self.ctx); + let previous_payload = self.ctx.codec_term_payloads.remove(&addr); + let new_payload = updater(previous_payload); + self.ctx.codec_term_payloads.insert(addr, new_payload); + } + fn index_document(&mut self, doc: &D) -> crate::Result<()> { let doc_id = self.max_doc; diff --git a/src/indexer/single_segment_index_writer.rs b/src/indexer/single_segment_index_writer.rs index a79bd7499..22153027f 100644 --- a/src/indexer/single_segment_index_writer.rs +++ b/src/indexer/single_segment_index_writer.rs @@ -1,3 +1,4 @@ +use std::any::Any; use std::marker::PhantomData; use crate::codec::StandardCodec; @@ -6,6 +7,7 @@ use crate::indexer::operation::AddOperation; use crate::indexer::segment_updater::save_metas; use crate::indexer::SegmentWriter; use crate::schema::document::Document; +use crate::schema::Field; use crate::{Directory, Index, IndexMeta, Opstamp, Segment, TantivyDocument}; #[doc(hidden)] @@ -42,6 +44,43 @@ impl SingleSegmentIndexWriter .add_document(AddOperation { opstamp, document }) } + /// Attaches or updates a codec-specific payload on a term of a regular + /// (non-JSON) field. + /// + /// `value_bytes` is the serialized term value, i.e. exactly what would be + /// appended after the field id (the raw text bytes for a str field, or the + /// big-endian bytes for a numeric field). + /// + /// The term does not need to belong to any document: if it does not exist + /// yet, it is created with an empty recorder so it still gets serialized. + /// `updater` receives the previously registered payload (`None` if absent) + /// and returns the payload to store. The payload is handed to the codec at + /// the beginning of the term during serialization. + pub fn update_term_payload( + &mut self, + field: Field, + value_bytes: &[u8], + updater: impl FnOnce(Option>) -> Box, + ) { + self.segment_writer + .update_term_payload(field, value_bytes, updater); + } + + /// Same as [`Self::update_term_payload`] for a JSON field. + /// + /// `value_bytes` must be the type-tagged value (`[type code][value]`), the + /// representation that follows the path within a JSON term. + pub fn update_json_term_payload( + &mut self, + field: Field, + json_path: &str, + value_bytes: &[u8], + updater: impl FnOnce(Option>) -> Box, + ) { + self.segment_writer + .update_json_term_payload(field, json_path, value_bytes, updater); + } + pub fn finalize(self) -> crate::Result> { let max_doc = self.segment_writer.max_doc(); self.segment_writer.finalize()?; @@ -60,3 +99,231 @@ impl SingleSegmentIndexWriter Ok(segment.index().clone()) } } + +#[cfg(test)] +mod tests { + use std::any::Any; + use std::cell::RefCell; + use std::io; + + use super::SingleSegmentIndexWriter; + use crate::codec::postings::{PostingsCodec, PostingsSerializer}; + use crate::codec::standard::postings::{ + SegmentPostings, StandardPostingsCodec, StandardPostingsSerializer, + }; + use crate::codec::Codec; + use crate::fieldnorm::FieldNormReader; + use crate::schema::{IndexRecordOption, Schema, Type, STRING}; + use crate::{DocId, Score, Term}; + + // The codec is round-tripped through `from_json_props` when the index is + // opened, so it cannot carry the capture sink itself. We use a thread-local + // sink instead: the `SingleSegmentIndexWriter` is single-threaded, so + // serialization runs on the test thread, and each test owns its own + // thread-local (clear it at the start of the test). + thread_local! { + static CAPTURED_PAYLOADS: RefCell> = const { RefCell::new(Vec::new()) }; + } + + fn reset_captured() { + CAPTURED_PAYLOADS.with(|captured| captured.borrow_mut().clear()); + } + + fn captured_payloads() -> Vec { + CAPTURED_PAYLOADS.with(|captured| captured.borrow().clone()) + } + + /// A postings serializer that delegates to the standard one, but records + /// the `u64` payload value of every term that carries a codec payload. + struct CapturingPostingsSerializer { + inner: StandardPostingsSerializer, + } + + impl PostingsSerializer for CapturingPostingsSerializer { + fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool) { + self.inner.new_term(term_doc_freq, record_term_freq); + } + + fn set_term_payload(&mut self, payload: &dyn Any) { + let value = *payload + .downcast_ref::() + .expect("payload should be a u64"); + CAPTURED_PAYLOADS.with(|captured| captured.borrow_mut().push(value)); + } + + fn write_doc(&mut self, doc_id: DocId, term_freq: u32) { + self.inner.write_doc(doc_id, term_freq); + } + + fn close_term(&mut self, doc_freq: u32, wrt: &mut impl io::Write) -> io::Result<()> { + self.inner.close_term(doc_freq, wrt) + } + } + + #[derive(Clone, Debug)] + struct CapturingPostingsCodec; + + impl PostingsCodec for CapturingPostingsCodec { + type PostingsSerializer = CapturingPostingsSerializer; + type Postings = SegmentPostings; + + fn new_serializer( + &self, + avg_fieldnorm: Score, + mode: IndexRecordOption, + fieldnorm_reader: Option, + ) -> Self::PostingsSerializer { + CapturingPostingsSerializer { + inner: StandardPostingsCodec.new_serializer(avg_fieldnorm, mode, fieldnorm_reader), + } + } + + fn load_postings( + &self, + doc_freq: u32, + postings_data: common::OwnedBytes, + record_option: IndexRecordOption, + requested_option: IndexRecordOption, + positions_data: Option, + ) -> io::Result { + StandardPostingsCodec.load_postings( + doc_freq, + postings_data, + record_option, + requested_option, + positions_data, + ) + } + } + + #[derive(Clone, Debug, Default)] + struct CapturingCodec; + + impl Codec for CapturingCodec { + type PostingsCodec = CapturingPostingsCodec; + + const ID: &'static str = "test-capturing-codec"; + + fn from_json_props(_json_value: &serde_json::Value) -> crate::Result { + Ok(CapturingCodec) + } + + fn to_json_props(&self) -> serde_json::Value { + serde_json::Value::Null + } + + fn postings_codec(&self) -> &Self::PostingsCodec { + &CapturingPostingsCodec + } + } + + fn build_writer(schema: Schema) -> SingleSegmentIndexWriter { + let index = crate::IndexBuilder::default() + .codec(CapturingCodec) + .schema(schema) + .create_in_ram() + .unwrap(); + SingleSegmentIndexWriter::new(index, 15_000_000).unwrap() + } + + #[test] + fn test_update_term_payload_regular_field() { + reset_captured(); + let mut schema_builder = Schema::builder(); + let text = schema_builder.add_text_field("text", STRING); + let schema = schema_builder.build(); + let mut writer = build_writer(schema); + + writer.add_document(crate::doc!(text => "alpha")).unwrap(); + writer.add_document(crate::doc!(text => "beta")).unwrap(); + writer.add_document(crate::doc!(text => "gamma")).unwrap(); + + // Existing term that belongs to a document. + writer.update_term_payload(text, b"beta", |previous_payload| { + assert!(previous_payload.is_none()); + Box::new(100u64) + }); + // Updating the same term: the previous payload is handed back. + writer.update_term_payload(text, b"beta", |previous_payload| { + let previous = previous_payload.expect("expected the previous payload"); + assert_eq!(*previous.downcast::().unwrap(), 100u64); + Box::new(101u64) + }); + // Brand-new term that belongs to no document: an empty recorder is + // created so it still lands in the term dictionary. + writer.update_term_payload(text, b"zeta", |previous_payload| { + assert!(previous_payload.is_none()); + Box::new(200u64) + }); + + let index = writer.finalize().unwrap(); + + // Terms are serialized in sorted order: alpha, beta, gamma, zeta. + // Only beta and zeta carry a payload. + assert_eq!(captured_payloads(), vec![101u64, 200u64]); + + let searcher = index.reader().unwrap().searcher(); + let segment_reader = searcher.segment_reader(0); + let inverted_index = segment_reader.inverted_index(text).unwrap(); + + let beta_info = inverted_index + .get_term_info(&Term::from_field_text(text, "beta")) + .unwrap() + .expect("beta should be in the dictionary"); + assert_eq!(beta_info.doc_freq, 1); + + let zeta_info = inverted_index + .get_term_info(&Term::from_field_text(text, "zeta")) + .unwrap() + .expect("zeta (no document) should still be in the dictionary"); + assert_eq!(zeta_info.doc_freq, 0); + } + + #[test] + fn test_update_json_term_payload() { + reset_captured(); + let mut schema_builder = Schema::builder(); + let json_field = schema_builder.add_json_field("json", STRING); + let schema = schema_builder.build(); + let mut writer = build_writer(schema); + + writer + .add_document(crate::doc!(json_field => serde_json::json!({"name": "hello"}))) + .unwrap(); + + let str_value = |value: &str| { + let mut bytes = vec![Type::Str.to_code()]; + bytes.extend_from_slice(value.as_bytes()); + bytes + }; + + // Existing str JSON term (path "name", value "hello"). + writer.update_json_term_payload(json_field, "name", &str_value("hello"), |previous| { + assert!(previous.is_none()); + Box::new(1u64) + }); + // Brand-new str JSON term with no document. + writer.update_json_term_payload(json_field, "name", &str_value("world"), |previous| { + assert!(previous.is_none()); + Box::new(2u64) + }); + // Brand-new non-str (numeric) JSON term with no document: exercises the + // DocIdRecorder branch of `ensure_term`. + let numeric_value = { + let mut bytes = vec![Type::I64.to_code()]; + bytes.extend_from_slice(&[0u8; 8]); + bytes + }; + writer.update_json_term_payload(json_field, "count", &numeric_value, |previous| { + assert!(previous.is_none()); + Box::new(3u64) + }); + + // Should not panic and should serialize cleanly. + let _index = writer.finalize().unwrap(); + + let mut got = captured_payloads(); + got.sort_unstable(); + assert_eq!(got, vec![1u64, 2u64, 3u64]); + } +} diff --git a/src/lib.rs b/src/lib.rs index efabe0352..4a60d65d6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -221,7 +221,7 @@ pub mod snippet; use std::fmt; pub use census::{Inventory, TrackedObject}; -pub use common::{f64_to_u64, i64_to_u64, u64_to_f64, u64_to_i64, HasLen}; +pub use common::{self, f64_to_u64, i64_to_u64, u64_to_f64, u64_to_i64, HasLen}; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; diff --git a/src/postings/indexing_context.rs b/src/postings/indexing_context.rs index 2675476f3..153b74957 100644 --- a/src/postings/indexing_context.rs +++ b/src/postings/indexing_context.rs @@ -1,4 +1,7 @@ -use stacker::{ArenaHashMap, MemoryArena}; +use std::any::Any; + +use fnv::FnvHashMap; +use stacker::{Addr, ArenaHashMap, MemoryArena}; use crate::indexer::path_to_unordered_id::PathToUnorderedId; @@ -11,6 +14,15 @@ pub(crate) struct IndexingContext { /// Arena is a memory arena that stores posting lists / term frequencies / positions. pub arena: MemoryArena, pub path_to_unordered_id: PathToUnorderedId, + /// Optional codec-specific payload attached to a term, keyed by the value + /// `Addr` of the term's recorder in `term_index`. + /// + /// Hidden contract: keying on `Addr` is sound because a term's recorder + /// address never changes once allocated (the arena only appends, and + /// `subscribe` updates the recorder in place). The payload is therefore + /// looked up by `Addr` at serialization time and fed to the codec's + /// postings serializer at the beginning of the term. + pub codec_term_payloads: FnvHashMap>, } impl IndexingContext { @@ -21,6 +33,7 @@ impl IndexingContext { arena: MemoryArena::default(), term_index, path_to_unordered_id: PathToUnorderedId::default(), + codec_term_payloads: FnvHashMap::default(), } } diff --git a/src/postings/json_postings_writer.rs b/src/postings/json_postings_writer.rs index 030508a8d..f44252085 100644 --- a/src/postings/json_postings_writer.rs +++ b/src/postings/json_postings_writer.rs @@ -96,6 +96,20 @@ impl PostingsWriter for JsonPostingsWriter { Ok(()) } + fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr { + // JSON term key layout: `[field:4][unordered_path_id:4][type code][value]`. + // Str values are recorded with `Rec`, all other types with `DocIdRecorder` + // (mirroring the dispatch in `serialize`). + let typ = Type::from_code(serialized_term[8]).expect("Invalid type code in JSON term"); + if typ == Type::Str { + ctx.term_index + .get_or_create_value_addr::(serialized_term, Rec::default) + } else { + ctx.term_index + .get_or_create_value_addr::(serialized_term, DocIdRecorder::default) + } + } + fn total_num_tokens(&self) -> u64 { self.str_posting_writer.total_num_tokens() + self.non_str_posting_writer.total_num_tokens() } diff --git a/src/postings/mod.rs b/src/postings/mod.rs index da42b7e44..13bafaf37 100644 --- a/src/postings/mod.rs +++ b/src/postings/mod.rs @@ -21,7 +21,8 @@ pub use postings::DocFreq; pub(crate) use stacker::compute_table_memory_size; pub(crate) use self::indexing_context::IndexingContext; -pub(crate) use self::per_field_postings_writer::PerFieldPostingsWriter; +#[doc(hidden)] +pub use self::per_field_postings_writer::PerFieldPostingsWriter; pub use self::postings::Postings; pub(crate) use self::postings_writer::{ serialize_postings, IndexingPosition, PostingsWriter, PostingsWriterEnum, diff --git a/src/postings/per_field_postings_writer.rs b/src/postings/per_field_postings_writer.rs index bf547cf5a..3ea6afa92 100644 --- a/src/postings/per_field_postings_writer.rs +++ b/src/postings/per_field_postings_writer.rs @@ -3,7 +3,7 @@ use crate::postings::postings_writer::{PostingsWriterEnum, SpecializedPostingsWr use crate::postings::recorder::{DocIdRecorder, TermFrequencyRecorder, TfAndPositionRecorder}; use crate::schema::{Field, FieldEntry, FieldType, IndexRecordOption, Schema}; -pub(crate) struct PerFieldPostingsWriter { +pub struct PerFieldPostingsWriter { per_field_postings_writers: Vec, } @@ -22,7 +22,7 @@ impl PerFieldPostingsWriter { &self.per_field_postings_writers[field.field_id() as usize] } - pub(crate) fn get_for_field_mut(&mut self, field: Field) -> &mut PostingsWriterEnum { + pub fn get_for_field_mut(&mut self, field: Field) -> &mut PostingsWriterEnum { &mut self.per_field_postings_writers[field.field_id() as usize] } } diff --git a/src/postings/postings_writer.rs b/src/postings/postings_writer.rs index c6f4466e6..a0b3c85c8 100644 --- a/src/postings/postings_writer.rs +++ b/src/postings/postings_writer.rs @@ -196,6 +196,21 @@ impl PostingsWriter for PostingsWriterEnum { } } + fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr { + match self { + PostingsWriterEnum::DocId(writer) => writer.ensure_term(serialized_term, ctx), + PostingsWriterEnum::DocIdTf(writer) => writer.ensure_term(serialized_term, ctx), + PostingsWriterEnum::DocTfAndPosition(writer) => { + writer.ensure_term(serialized_term, ctx) + } + PostingsWriterEnum::JsonDocId(writer) => writer.ensure_term(serialized_term, ctx), + PostingsWriterEnum::JsonDocIdTf(writer) => writer.ensure_term(serialized_term, ctx), + PostingsWriterEnum::JsonDocTfAndPosition(writer) => { + writer.ensure_term(serialized_term, ctx) + } + } + } + /// Tokenize a text and subscribe all of its token. fn index_text( &mut self, @@ -263,6 +278,15 @@ pub(crate) trait PostingsWriter: Send + Sync { serializer: &mut FieldSerializer, ) -> io::Result<()>; + /// Ensures `serialized_term` has an entry in the term index, creating an + /// empty recorder (matching this writer's indexing option) if the term is + /// not present yet, and returns the value `Addr` of its recorder. + /// + /// An existing recorder is never overwritten, so the term keeps any + /// posting data already recorded for it. This is used to attach a + /// codec-specific payload to a term that may belong to no document. + fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr; + /// Tokenize a text and subscribe all of its token. fn index_text( &mut self, @@ -322,6 +346,10 @@ impl SpecializedPostingsWriter { let recorder: Rec = ctx.term_index.read(addr); let term_doc_freq = recorder.term_doc_freq().unwrap_or(0u32); serializer.new_term(term, term_doc_freq, recorder.has_term_freq())?; + if let Some(payload) = ctx.codec_term_payloads.get(&addr) { + // `&(dyn Any + Send)` upcasts to `&dyn Any`. + serializer.set_term_payload(payload.as_ref()); + } recorder.serialize(&ctx.arena, serializer, buffer_lender); serializer.close_term()?; Ok(()) @@ -372,6 +400,11 @@ impl PostingsWriter for SpecializedPostingsWriter { Ok(()) } + fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr { + ctx.term_index + .get_or_create_value_addr::(serialized_term, Rec::default) + } + fn total_num_tokens(&self) -> u64 { self.total_num_tokens } diff --git a/src/postings/serializer.rs b/src/postings/serializer.rs index 63fc5c0d8..2e8386187 100644 --- a/src/postings/serializer.rs +++ b/src/postings/serializer.rs @@ -203,6 +203,13 @@ impl<'a, C: Codec> FieldSerializer<'a, C> { self.new_term(term, 0, false) } + /// Forwards a codec-specific per-term payload to the postings serializer. + /// + /// Must be called after `new_term` and before any `write_doc`. + pub fn set_term_payload(&mut self, payload: &dyn std::any::Any) { + self.postings_serializer.set_term_payload(payload); + } + /// Serialize the information that a document contains for the current term: /// its term frequency, and the position deltas. /// diff --git a/src/query/boolean_query/boolean_weight.rs b/src/query/boolean_query/boolean_weight.rs index 8088fa233..14e559edb 100644 --- a/src/query/boolean_query/boolean_weight.rs +++ b/src/query/boolean_query/boolean_weight.rs @@ -279,7 +279,6 @@ impl BooleanWeight { Some(exclude_scorers_union) }; - let include_scorer = match (should_scorers, must_scorers) { (ShouldScorersCombinationMethod::Ignored, must_scorers) => { // No SHOULD clauses (or they were absorbed into MUST). diff --git a/src/query/boolean_query/mod.rs b/src/query/boolean_query/mod.rs index 71f56d95d..c001cce94 100644 --- a/src/query/boolean_query/mod.rs +++ b/src/query/boolean_query/mod.rs @@ -1,7 +1,6 @@ mod boolean_query; mod boolean_weight; - pub use self::boolean_query::BooleanQuery; pub use self::boolean_weight::BooleanWeight; diff --git a/stacker/src/arena_hashmap.rs b/stacker/src/arena_hashmap.rs index 9f3dfbbc5..244c30f83 100644 --- a/stacker/src/arena_hashmap.rs +++ b/stacker/src/arena_hashmap.rs @@ -83,6 +83,28 @@ impl ArenaHashMap { self.shared_arena_hashmap .mutate_or_create(key, &mut self.memory_arena, updater); } + + /// Returns the address of the value associated to `key`, creating an entry + /// with `make_default()` if the key is not present yet (an existing value + /// is left untouched). + /// + /// See [`SharedArenaHashMap::get_or_create_value_addr`] for the address + /// stability guarantees. + #[inline] + pub fn get_or_create_value_addr( + &mut self, + key: &[u8], + make_default: impl FnOnce() -> V, + ) -> Addr + where + V: Copy + 'static, + { + self.shared_arena_hashmap.get_or_create_value_addr( + key, + &mut self.memory_arena, + make_default, + ) + } } #[cfg(test)] @@ -121,6 +143,28 @@ mod tests { assert_eq!(hash_map.get::(b"abc"), None); } + #[test] + fn test_get_or_create_value_addr() { + let mut hash_map: ArenaHashMap = ArenaHashMap::default(); + // Creates the entry with the default value. + let addr_abc = hash_map.get_or_create_value_addr(b"abc", || 7u32); + assert_eq!(hash_map.read::(addr_abc), 7u32); + // Returns the same address and does NOT overwrite an existing value. + let addr_abc_again = hash_map.get_or_create_value_addr(b"abc", || 99u32); + assert_eq!(addr_abc_again, addr_abc); + assert_eq!(hash_map.read::(addr_abc), 7u32); + // A different key gets its own entry. + let addr_def = hash_map.get_or_create_value_addr(b"def", || 5u32); + assert_ne!(addr_def, addr_abc); + assert_eq!(hash_map.read::(addr_def), 5u32); + // The address matches the one yielded by `iter`. + for (key, addr) in hash_map.iter() { + if key == b"abc" { + assert_eq!(addr, addr_abc); + } + } + } + #[test] fn test_many_terms() { let mut terms: Vec = (0..20_000).map(|val| val.to_string()).collect(); diff --git a/stacker/src/memory_arena.rs b/stacker/src/memory_arena.rs index 383b96438..6670718cb 100644 --- a/stacker/src/memory_arena.rs +++ b/stacker/src/memory_arena.rs @@ -36,7 +36,7 @@ const PAGE_SIZE: usize = 1 << NUM_BITS_PAGE_ADDR; // pages are 1 MB large /// page of memory. /// /// The last 20 bits are an address within this page of memory. -#[derive(Copy, Clone, Debug)] +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] pub struct Addr(u32); impl Addr { diff --git a/stacker/src/shared_arena_hashmap.rs b/stacker/src/shared_arena_hashmap.rs index 6964272f0..66a5cdf37 100644 --- a/stacker/src/shared_arena_hashmap.rs +++ b/stacker/src/shared_arena_hashmap.rs @@ -347,6 +347,67 @@ impl SharedArenaHashMap { kv = self.table[bucket]; } } + + /// Returns the address of the value associated to `key`. + /// + /// If the key is not present yet, a new entry is created with the value + /// returned by `make_default()`. If the key is already present, the stored + /// value is left untouched and its address is returned. + /// + /// The returned `Addr` is the value address, i.e. the same address yielded + /// by [`Self::iter`] and consumed by [`MemoryArena::read`]. It remains valid + /// for the lifetime of the arena: arena allocations only ever append, and a + /// table resize relocates buckets, not the arena-backed key/value data. + /// + /// The key will be truncated to `u16::MAX` bytes. + #[inline] + pub fn get_or_create_value_addr( + &mut self, + key: &[u8], + memory_arena: &mut MemoryArena, + make_default: impl FnOnce() -> V, + ) -> Addr + where + V: Copy + 'static, + { + if self.is_saturated() { + self.resize(); + } + // Limit the key size to u16::MAX + let key = &key[..std::cmp::min(key.len(), u16::MAX as usize)]; + let hash = self.get_hash(key); + let mut probe = self.probe(hash); + let mut bucket = probe.next_probe(); + let mut kv: KeyValue = self.table[bucket]; + loop { + if kv.is_empty() { + // The key does not exist yet: create it with the default value. + let val = make_default(); + let num_bytes = std::mem::size_of::() + key.len() + std::mem::size_of::(); + let key_addr = memory_arena.allocate_space(num_bytes); + { + let data = memory_arena.slice_mut(key_addr, num_bytes); + let key_len_bytes: [u8; 2] = (key.len() as u16).to_le_bytes(); + data[..2].copy_from_slice(&key_len_bytes); + let stop = 2 + key.len(); + fast_short_slice_copy(key, &mut data[2..stop]); + store(&mut data[stop..], val); + } + self.set_bucket(hash, key_addr, bucket); + return key_addr.offset(2 + key.len() as u32); + } + if kv.hash == hash + && let Some(val_addr) = + self.get_value_addr_if_key_match(key, kv.key_value_addr, memory_arena) + { + // The key already exists: leave its value untouched. + return val_addr; + } + // This allows fetching the next bucket before the loop jmp + bucket = probe.next_probe(); + kv = self.table[bucket]; + } + } } #[cfg(test)]