use super::{EnumVariant, Field, FieldId, FieldTy, Model, ModelId, VariantId};
use crate::{Result, stmt};
use indexmap::IndexMap;
#[derive(Debug)]
pub enum Resolved<'a> {
Field(&'a Field),
Variant(&'a EnumVariant),
}
#[derive(Debug, Default)]
pub struct Schema {
pub models: IndexMap<ModelId, Model>,
}
#[derive(Default)]
struct Builder {
models: IndexMap<ModelId, Model>,
}
impl Schema {
pub fn from_macro(models: impl IntoIterator<Item = Model>) -> Result<Self> {
Builder::from_macro(models)
}
pub fn field(&self, id: FieldId) -> &Field {
let fields = match self.model(id.model) {
Model::Root(root) => &root.fields,
Model::EmbeddedStruct(embedded) => &embedded.fields,
Model::EmbeddedEnum(e) => &e.fields,
};
fields.get(id.index).expect("invalid field ID")
}
pub fn variant(&self, id: VariantId) -> &EnumVariant {
let Model::EmbeddedEnum(e) = self.model(id.model) else {
panic!("VariantId references a non-enum model");
};
e.variants.get(id.index).expect("invalid variant index")
}
pub fn models(&self) -> impl Iterator<Item = &Model> {
self.models.values()
}
pub fn get_model(&self, id: impl Into<ModelId>) -> Option<&Model> {
self.models.get(&id.into())
}
pub fn model(&self, id: impl Into<ModelId>) -> &Model {
self.models.get(&id.into()).expect("invalid model ID")
}
pub fn resolve<'a>(
&'a self,
root: &'a Model,
projection: &stmt::Projection,
) -> Option<Resolved<'a>> {
let [first, rest @ ..] = projection.as_slice() else {
return None;
};
let mut current_field = root.as_root_unwrap().fields.get(*first)?;
let mut steps = rest.iter();
while let Some(step) = steps.next() {
match ¤t_field.ty {
FieldTy::Primitive(..) => {
return None;
}
FieldTy::Embedded(embedded) => {
let target = self.model(embedded.target);
match target {
Model::EmbeddedStruct(s) => {
current_field = s.fields.get(*step)?;
}
Model::EmbeddedEnum(e) => {
let variant = e.variants.get(*step)?;
if let Some(field_step) = steps.next() {
current_field = e.fields.get(*field_step)?;
} else {
return Some(Resolved::Variant(variant));
}
}
_ => return None,
}
}
FieldTy::BelongsTo(belongs_to) => {
current_field = belongs_to.target(self).as_root_unwrap().fields.get(*step)?;
}
FieldTy::HasMany(has_many) => {
current_field = has_many.target(self).as_root_unwrap().fields.get(*step)?;
}
FieldTy::HasOne(has_one) => {
current_field = has_one.target(self).as_root_unwrap().fields.get(*step)?;
}
};
}
Some(Resolved::Field(current_field))
}
pub fn resolve_field<'a>(
&'a self,
root: &'a Model,
projection: &stmt::Projection,
) -> Option<&'a Field> {
match self.resolve(root, projection) {
Some(Resolved::Field(field)) => Some(field),
_ => None,
}
}
pub fn resolve_field_path<'a>(&'a self, path: &stmt::Path) -> Option<&'a Field> {
let model = self.model(path.root.as_model_unwrap());
self.resolve_field(model, &path.projection)
}
}
impl Builder {
pub(crate) fn from_macro(models: impl IntoIterator<Item = Model>) -> Result<Schema> {
let mut builder = Self { ..Self::default() };
for model in models {
builder.models.insert(model.id(), model);
}
builder.process_models()?;
builder.into_schema()
}
fn into_schema(self) -> Result<Schema> {
Ok(Schema {
models: self.models,
})
}
fn process_models(&mut self) -> Result<()> {
self.link_relations()?;
Ok(())
}
fn link_relations(&mut self) -> crate::Result<()> {
for curr in 0..self.models.len() {
if self.models[curr].is_embedded() {
continue;
}
for index in 0..self.models[curr].as_root_unwrap().fields.len() {
let model = &self.models[curr];
let src = model.id();
let field = &model.as_root_unwrap().fields[index];
if let FieldTy::HasMany(has_many) = &field.ty {
let target = has_many.target;
let field_name = field.name.app_name.clone();
let pair = self.find_has_many_pair(src, target, &field_name)?;
self.models[curr].as_root_mut_unwrap().fields[index]
.ty
.as_has_many_mut_unwrap()
.pair = pair;
}
}
}
for curr in 0..self.models.len() {
if self.models[curr].is_embedded() {
continue;
}
for index in 0..self.models[curr].as_root_unwrap().fields.len() {
let model = &self.models[curr];
let src = model.id();
let field = &model.as_root_unwrap().fields[index];
match &field.ty {
FieldTy::HasOne(has_one) => {
let target = has_one.target;
let field_name = field.name.app_name.clone();
let pair = match self.find_belongs_to_pair(src, target, &field_name)? {
Some(pair) => pair,
None => {
return Err(crate::Error::invalid_schema(format!(
"field `{}::{}` has no matching `BelongsTo` relation on the target model",
self.models[curr].name().upper_camel_case(),
field_name,
)));
}
};
self.models[curr].as_root_mut_unwrap().fields[index]
.ty
.as_has_one_mut_unwrap()
.pair = pair;
}
FieldTy::BelongsTo(belongs_to) => {
assert!(!belongs_to.foreign_key.is_placeholder());
continue;
}
_ => {}
}
}
}
for curr in 0..self.models.len() {
if self.models[curr].is_embedded() {
continue;
}
for index in 0..self.models[curr].as_root_unwrap().fields.len() {
let model = &self.models[curr];
let field_id = model.as_root_unwrap().fields[index].id;
let pair = match &self.models[curr].as_root_unwrap().fields[index].ty {
FieldTy::BelongsTo(belongs_to) => {
let mut pair = None;
let target = match self.models.get_index_of(&belongs_to.target) {
Some(target) => target,
None => {
let model = &self.models[curr];
return Err(crate::Error::invalid_schema(format!(
"field `{}::{}` references a model that was not registered \
with the schema; did you forget to register it with `Db::builder()`?",
model.name().upper_camel_case(),
model.as_root_unwrap().fields[index].name.app_name,
)));
}
};
for target_index in 0..self.models[target].as_root_unwrap().fields.len() {
pair = match &self.models[target].as_root_unwrap().fields[target_index]
.ty
{
FieldTy::HasMany(has_many) if has_many.pair == field_id => {
assert!(pair.is_none());
Some(
self.models[target].as_root_unwrap().fields[target_index]
.id,
)
}
FieldTy::HasOne(has_one) if has_one.pair == field_id => {
assert!(pair.is_none());
Some(
self.models[target].as_root_unwrap().fields[target_index]
.id,
)
}
_ => continue,
}
}
if pair.is_none() {
continue;
}
pair
}
_ => continue,
};
self.models[curr].as_root_mut_unwrap().fields[index]
.ty
.as_belongs_to_mut_unwrap()
.pair = pair;
}
}
Ok(())
}
fn find_belongs_to_pair(
&self,
src: ModelId,
target: ModelId,
field_name: &str,
) -> crate::Result<Option<FieldId>> {
let src_model = &self.models[&src];
let target = match self.models.get(&target) {
Some(target) => target,
None => {
return Err(crate::Error::invalid_schema(format!(
"field `{}::{}` references a model that was not registered with the schema; \
did you forget to register it with `Db::builder()`?",
src_model.name().upper_camel_case(),
field_name,
)));
}
};
let belongs_to: Vec<_> = target
.as_root_unwrap()
.fields
.iter()
.filter(|field| match &field.ty {
FieldTy::BelongsTo(rel) => rel.target == src,
_ => false,
})
.collect();
match &belongs_to[..] {
[field] => Ok(Some(field.id)),
[] => Ok(None),
_ => Err(crate::Error::invalid_schema(format!(
"model `{}` has more than one `BelongsTo` relation targeting `{}`",
target.name().upper_camel_case(),
src_model.name().upper_camel_case(),
))),
}
}
fn find_has_many_pair(
&mut self,
src: ModelId,
target: ModelId,
field_name: &str,
) -> crate::Result<FieldId> {
if let Some(field_id) = self.find_belongs_to_pair(src, target, field_name)? {
return Ok(field_id);
}
Err(crate::Error::invalid_schema(format!(
"field `{}::{}` has no matching `BelongsTo` relation on the target model",
self.models[&src].name().upper_camel_case(),
field_name,
)))
}
}