use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use memstead_schema::{Filterable, Schema};
use tantivy::schema::{
Field, STORED, STRING, Schema as TantivySchema, SchemaBuilder, TextFieldIndexing, TextOptions,
};
use super::tokenizer::MEMSTEAD_TOKENIZER;
#[derive(Clone, Debug)]
pub struct IndexFields {
pub schema: TantivySchema,
pub id: Field,
pub mem: Field,
pub entity_type: Field,
pub title: Field,
pub sections: BTreeMap<String, Field>,
pub metadata: BTreeMap<String, Field>,
}
impl IndexFields {
pub fn build(mem_schema: Option<&Arc<Schema>>) -> Self {
let mut builder = SchemaBuilder::new();
let id = builder.add_text_field("id", STRING | STORED);
let mem = builder.add_text_field("mem", STRING | STORED);
let entity_type = builder.add_text_field("entity_type", STRING | STORED);
let title = builder.add_text_field("title", text_options());
let section_keys: BTreeSet<String> = match mem_schema {
Some(schema) => schema
.types
.values()
.flat_map(|t| t.sections.iter().map(|s| s.key.clone()))
.collect(),
None => BTreeSet::new(),
};
let mut sections = BTreeMap::new();
for key in section_keys {
let f = builder.add_text_field(&format!("section_{key}"), text_options());
sections.insert(key, f);
}
let filterable_keys: BTreeSet<String> = match mem_schema {
Some(schema) => schema
.types
.values()
.flat_map(|t| t.metadata_fields.iter())
.filter(|f| matches!(f.filterable, Filterable::Equality | Filterable::Range))
.map(|f| f.key.clone())
.collect(),
None => BTreeSet::new(),
};
let mut metadata = BTreeMap::new();
for key in filterable_keys {
let f = builder.add_text_field(&format!("meta_{key}"), STRING);
metadata.insert(key, f);
}
let schema = builder.build();
Self {
schema,
id,
mem,
entity_type,
title,
sections,
metadata,
}
}
}
fn text_options() -> TextOptions {
TextOptions::default().set_indexing_options(
TextFieldIndexing::default()
.set_tokenizer(MEMSTEAD_TOKENIZER)
.set_index_option(tantivy::schema::IndexRecordOption::WithFreqsAndPositions),
)
}
#[cfg(test)]
mod tests {
use super::*;
use memstead_schema::Schema;
#[test]
fn build_without_schema_still_has_fixed_fields() {
let fields = IndexFields::build(None);
assert!(fields.schema.get_field("id").is_ok());
assert!(fields.schema.get_field("mem").is_ok());
assert!(fields.schema.get_field("title").is_ok());
assert!(fields.sections.is_empty());
assert!(fields.metadata.is_empty());
}
#[test]
fn build_with_default_schema_emits_section_fields() {
let schema = Schema::builtin_default();
let fields = IndexFields::build(Some(&schema));
assert!(fields.sections.contains_key("identity"));
assert!(fields.sections.contains_key("purpose"));
assert!(fields.schema.get_field("section_identity").is_ok());
}
#[test]
fn build_with_default_schema_emits_filterable_metadata() {
let schema = Schema::builtin_default();
let fields = IndexFields::build(Some(&schema));
assert!(
!fields.metadata.is_empty(),
"default schema should expose at least one filterable metadata field"
);
}
}