1use crate::{
2 node::{DataStore, Entity, Schema, SchemaNode},
3 prelude::*,
4};
5use std::collections::{BTreeMap, BTreeSet};
6
7pub struct EntityInfo {
13 name: String,
14 canister: String,
15}
16
17pub struct RelationEdge {
23 source_entity: String,
24 target_entity: String,
25 field: String,
26}
27
28pub fn validate_same_canister_relations(schema: &Schema, errs: &mut ErrorTree) {
31 let mut edges = Vec::new();
33 for (entity_path, entity) in schema.get_nodes::<Entity>() {
34 collect_entity_relations(schema, entity_path, entity, &mut edges);
35 }
36
37 let entity_info = build_entity_info_map(schema, errs);
39 for edge in edges {
40 let Some(source) = entity_info.get(&edge.source_entity) else {
41 continue;
42 };
43 let Some(target) = entity_info.get(&edge.target_entity) else {
44 continue;
45 };
46 if source.canister != target.canister {
47 err!(
48 errs,
49 "entity '{0}' (canister '{1}'), field '{2}', has a relation to entity '{3}' (canister '{4}'), which is not allowed",
50 source.name,
51 source.canister,
52 edge.field,
53 target.name,
54 target.canister
55 );
56 }
57 }
58}
59
60fn build_entity_info_map(schema: &Schema, errs: &mut ErrorTree) -> BTreeMap<String, EntityInfo> {
62 let mut entity_info = BTreeMap::new();
63 for (entity_path, entity) in schema.get_nodes::<Entity>() {
64 let store = match schema.cast_node::<DataStore>(entity.store) {
65 Ok(store) => store,
66 Err(e) => {
67 errs.add(e);
68 continue;
69 }
70 };
71
72 entity_info.insert(
73 entity_path.to_string(),
74 EntityInfo {
75 name: entity.resolved_name().to_string(),
76 canister: store.canister.to_string(),
77 },
78 );
79 }
80
81 entity_info
82}
83
84fn collect_entity_relations(
86 schema: &Schema,
87 entity_path: &str,
88 entity: &Entity,
89 edges: &mut Vec<RelationEdge>,
90) {
91 let mut visiting = BTreeSet::new();
92
93 for field in entity.fields.fields {
94 let mut field_path = vec![field.ident.to_string()];
95 collect_value_relations(
96 schema,
97 entity_path,
98 &field.value,
99 &mut field_path,
100 &mut visiting,
101 edges,
102 );
103 }
104}
105
106fn collect_value_relations(
108 schema: &Schema,
109 entity_path: &str,
110 value: &Value,
111 field_path: &mut Vec<String>,
112 visiting: &mut BTreeSet<String>,
113 edges: &mut Vec<RelationEdge>,
114) {
115 collect_item_relations(
116 schema,
117 entity_path,
118 &value.item,
119 field_path,
120 visiting,
121 edges,
122 );
123}
124
125fn collect_item_relations(
127 schema: &Schema,
128 entity_path: &str,
129 item: &Item,
130 field_path: &mut Vec<String>,
131 visiting: &mut BTreeSet<String>,
132 edges: &mut Vec<RelationEdge>,
133) {
134 if let Some(relation) = item.relation {
135 edges.push(RelationEdge {
136 source_entity: entity_path.to_string(),
137 target_entity: relation.to_string(),
138 field: format_field_path(field_path),
139 });
140 }
141
142 if let ItemTarget::Is(path) = &item.target {
143 traverse_type_node(schema, entity_path, path, field_path, visiting, edges);
144 }
145}
146
147fn traverse_type_node(
149 schema: &Schema,
150 entity_path: &str,
151 type_path: &str,
152 field_path: &mut Vec<String>,
153 visiting: &mut BTreeSet<String>,
154 edges: &mut Vec<RelationEdge>,
155) {
156 if !visiting.insert(type_path.to_string()) {
157 return;
158 }
159
160 let Some(node) = schema.get_node(type_path) else {
161 visiting.remove(type_path);
162 return;
163 };
164
165 match node {
166 SchemaNode::Record(record) => {
167 for field in record.fields.fields {
168 field_path.push(field.ident.to_string());
169 collect_value_relations(
170 schema,
171 entity_path,
172 &field.value,
173 field_path,
174 visiting,
175 edges,
176 );
177 field_path.pop();
178 }
179 }
180 SchemaNode::Enum(enumeration) => {
181 for variant in enumeration.variants {
182 let Some(value) = &variant.value else {
183 continue;
184 };
185 field_path.push(variant.ident.to_string());
186 collect_value_relations(schema, entity_path, value, field_path, visiting, edges);
187 field_path.pop();
188 }
189 }
190 SchemaNode::Tuple(tuple) => {
191 for (index, value) in tuple.values.iter().enumerate() {
192 field_path.push(format!("[{index}]"));
193 collect_value_relations(schema, entity_path, value, field_path, visiting, edges);
194 field_path.pop();
195 }
196 }
197 SchemaNode::List(list) => {
198 field_path.push("item".to_string());
199 collect_item_relations(schema, entity_path, &list.item, field_path, visiting, edges);
200 field_path.pop();
201 }
202 SchemaNode::Set(set) => {
203 field_path.push("item".to_string());
204 collect_item_relations(schema, entity_path, &set.item, field_path, visiting, edges);
205 field_path.pop();
206 }
207 SchemaNode::Map(map) => {
208 field_path.push("key".to_string());
209 collect_item_relations(schema, entity_path, &map.key, field_path, visiting, edges);
210 field_path.pop();
211
212 field_path.push("value".to_string());
213 collect_value_relations(schema, entity_path, &map.value, field_path, visiting, edges);
214 field_path.pop();
215 }
216 SchemaNode::Newtype(newtype) => {
217 field_path.push("item".to_string());
218 collect_item_relations(
219 schema,
220 entity_path,
221 &newtype.item,
222 field_path,
223 visiting,
224 edges,
225 );
226 field_path.pop();
227 }
228 _ => {}
229 }
230
231 visiting.remove(type_path);
232}
233
234fn format_field_path(field_path: &[String]) -> String {
236 field_path.join(".")
237}