diff --git a/crates/milli/src/update/new/document.rs b/crates/milli/src/update/new/document.rs index b1a2218f2..b8a78efc3 100644 --- a/crates/milli/src/update/new/document.rs +++ b/crates/milli/src/update/new/document.rs @@ -203,6 +203,83 @@ impl<'a, 'doc, 't, Mapper: FieldIdMapper> MergedDocument<'a, 'doc, 't, Mapper> { pub fn without_db(new_doc: DocumentFromVersions<'a, 'doc>) -> Self { Self { new_doc, db: None } } + + pub fn iter_merged_top_level_fields<'d>( + &self, + ) -> impl Iterator)>> + '_ + where + 't: 'd, + 'doc: 'd, + { + match &self.db { + Some(db) => { + let mut new_doc_it = self.new_doc.iter_top_level_fields(); + let mut db_it = db.iter_top_level_fields(); + let mut seen_fields = BTreeSet::new(); + + Either::Left(std::iter::from_fn(move || { + if let Some(next) = new_doc_it.next() { + let (name, updated_value) = match next { + Ok((name, updated_value)) => (name, updated_value), + Err(err) => return Some(Err(err)), + }; + seen_fields.insert(name); + let current = match db.top_level_field(name) { + Ok(current) => current, + Err(err) => return Some(Err(err)), + }; + + match current { + Some(current) => { + return Some(Ok(( + name, + MergedValue::CurrentAndUpdated(current, updated_value), + ))) + } + None => return Some(Ok((name, MergedValue::Updated(updated_value)))), + } + } + loop { + match db_it.next()? { + Ok((name, value)) => { + if seen_fields.contains(name) { + continue; + } + return Some(Ok((name, MergedValue::Current(value)))); + } + Err(err) => return Some(Err(err)), + } + } + })) + } + None => Either::Right(self.new_doc.iter_top_level_fields().map(|res| { + let (k, v) = res?; + Ok((k, MergedValue::Updated(v))) + })), + } + } + + pub fn merged_geo_field(&self) -> Result>> { + let db_geo_field = match self.db { + Some(db) => db.geo_field()?, + None => None, + }; + + let new_doc_geo_field = self.new_doc.geo_field()?; + + Ok(match (db_geo_field, new_doc_geo_field) { + (None, None) => None, + (None, Some(new_doc)) => Some(MergedValue::Updated(new_doc)), + (Some(db), None) => Some(MergedValue::Current(db)), + (Some(db), Some(new_doc)) => Some(MergedValue::CurrentAndUpdated(db, new_doc)), + }) + } +} + +pub enum MergedValue<'t, 'doc> { + Current(&'t RawValue), + Updated(&'doc RawValue), + CurrentAndUpdated(&'t RawValue, &'doc RawValue), } impl<'d, 'doc: 'd, 't: 'd, Mapper: FieldIdMapper> Document<'d>