1use std::{
49 collections::{HashMap, HashSet},
50 path::PathBuf,
51 sync::{Arc, Mutex},
52};
53
54use sha2::{Digest, Sha256};
55use sqlx::PgPool;
56use uuid::Uuid;
57
58use crate::{
59 auth::Actor,
60 error::{AppError, AppResult},
61 invplace::Override,
62 mesh::{Mesh, MAX_TRIANGLES},
63 nxasm::{self, Kind},
64 preview::{PreviewFormat, MAX_PARALLEL},
65 services::access,
66 AppState,
67};
68
69pub const MAX_DEPTH: usize = 8;
71pub const MAX_PLACED: usize = 50_000;
73
74#[derive(Clone, Debug, sqlx::FromRow)]
76struct Found {
77 key: String,
78 sha256: String,
79 format: String,
80 revision_id: Uuid,
81}
82
83#[derive(Clone, Copy, Debug, PartialEq, Eq)]
86enum Family {
87 Nx,
88 SolidWorks,
89 Inventor,
90 Catia,
91 Creo,
92}
93
94impl Family {
95 fn of(format: &str) -> Option<Family> {
96 match format {
97 "nx-part" => Some(Family::Nx),
98 "solidworks-assembly" => Some(Family::SolidWorks),
99 "inventor-assembly" => Some(Family::Inventor),
100 "catia-product" => Some(Family::Catia),
101 "creo-assembly" => Some(Family::Creo),
102 _ => None,
103 }
104 }
105
106 fn key(self, component: &nxasm::Component) -> String {
109 let name = &component.name;
110 match self {
111 Family::Nx => format!("{}.prt", name.to_lowercase()),
112 Family::SolidWorks | Family::Inventor | Family::Creo => name.to_lowercase(),
113 Family::Catia => component
114 .document
115 .clone()
116 .unwrap_or_else(|| name.to_lowercase()),
117 }
118 }
119
120 fn accepts(self, format: &str) -> bool {
121 match self {
122 Family::Nx => format == "nx-part",
123 Family::SolidWorks => matches!(format, "solidworks-part" | "solidworks-assembly"),
124 Family::Inventor => matches!(format, "inventor-part" | "inventor-assembly"),
125 Family::Catia => matches!(format, "catia-part" | "catia-product"),
126 Family::Creo => matches!(format, "creo-part" | "creo-assembly"),
127 }
128 }
129
130 fn part_format(self) -> PreviewFormat {
132 match self {
133 Family::Nx => PreviewFormat::NxPart,
134 Family::SolidWorks => PreviewFormat::SwPart,
135 Family::Inventor => PreviewFormat::InvPart,
136 Family::Catia => PreviewFormat::CatiaPart,
137 Family::Creo => PreviewFormat::CreoPart,
138 }
139 }
140}
141
142#[derive(Debug)]
144pub struct Plan {
145 pub parts: Vec<(String, nxasm::Component)>,
147 pub missing: Vec<String>,
149 pub key: String,
151 pub part_format: PreviewFormat,
153}
154
155async fn resolve(
157 pool: &PgPool,
158 actor: &Actor,
159 revision_id: Uuid,
160 keys: &[String],
161) -> AppResult<HashMap<String, Found>> {
162 let keys: Vec<String> = keys
163 .iter()
164 .cloned()
165 .collect::<std::collections::BTreeSet<_>>()
166 .into_iter()
167 .collect();
168 let rows = sqlx::query_as::<_, Found>(&format!(
169 "SELECT lower(f.filename) AS key, f.sha256, COALESCE(f.format, '') AS format,
170 f.revision_id
171 FROM files f
172 JOIN revisions r ON r.id = f.revision_id
173 JOIN items i ON i.id = r.item_id
174 WHERE lower(f.filename) = ANY($1)
175 AND (r.id = $2 OR (r.item_id <> (SELECT item_id FROM revisions WHERE id = $2)
176 AND r.seq = (SELECT MAX(x.seq) FROM revisions x WHERE x.item_id = r.item_id)))
177 AND {}
178 ORDER BY (r.id = $2) DESC, i.number, f.filename",
179 access::item_visible("i", 3)
180 ))
181 .bind(&keys)
182 .bind(revision_id)
183 .bind(&actor.hidden_projects)
184 .fetch_all(pool)
185 .await?;
186 let mut out = HashMap::new();
187 for row in rows {
188 out.entry(row.key.clone()).or_insert(row);
190 }
191 Ok(out)
192}
193
194const MAX_OWN_CATIA: i64 = 2_000;
196
197async fn resolve_catia(
203 st: &AppState,
204 actor: &Actor,
205 revision_id: Uuid,
206 components: &[nxasm::Component],
207) -> AppResult<HashMap<String, Found>> {
208 let names = |c: &nxasm::Component| {
209 let n = c.name.to_lowercase();
210 [format!("{n}.catpart"), format!("{n}.catproduct")]
211 };
212 let keys: Vec<String> = components.iter().flat_map(names).collect();
213 let by_name = resolve(&st.db, actor, revision_id, &keys).await?;
214 let mut candidates = sqlx::query_as::<_, Found>(&format!(
215 "SELECT lower(f.filename) AS key, f.sha256, COALESCE(f.format, '') AS format,
216 f.revision_id
217 FROM files f
218 JOIN revisions r ON r.id = f.revision_id
219 JOIN items i ON i.id = r.item_id
220 WHERE f.revision_id = $1 AND f.format IN ('catia-part', 'catia-product')
221 AND {}
222 ORDER BY f.filename
223 LIMIT {MAX_OWN_CATIA}",
224 access::item_visible("i", 2)
225 ))
226 .bind(revision_id)
227 .bind(&actor.hidden_projects)
228 .fetch_all(&st.db)
229 .await?;
230 candidates.extend(by_name.values().cloned());
231 let mut documents: HashMap<String, Found> = HashMap::new();
233 let mut read: HashSet<String> = HashSet::new();
234 for file in candidates {
235 if !read.insert(file.sha256.clone()) {
236 continue;
237 }
238 let path = st.vault.object_path(&file.sha256)?;
239 let document = tokio::task::spawn_blocking(move || crate::catprod::document_of(&path))
240 .await
241 .map_err(|e| AppError::Other(anyhow::anyhow!("reading the product failed: {e}")))?;
242 if let Ok(Some(document)) = document {
243 documents.entry(document).or_insert(file);
244 }
245 }
246 let mut out = HashMap::new();
247 for c in components {
248 let key = Family::Catia.key(c);
249 let by_document = c.document.as_ref().and_then(|d| documents.get(d));
250 let [part, product] = names(c);
251 if let Some(file) = by_document
252 .or_else(|| by_name.get(&part))
253 .or_else(|| by_name.get(&product))
254 {
255 out.insert(key, file.clone());
256 }
257 }
258 Ok(out)
259}
260
261const MAX_TREE: u64 = 64 * 1024 * 1024;
263
264async fn kind(st: &AppState, sha256: &str, format: &str) -> AppResult<Kind> {
269 let path = st.vault.object_path(sha256)?;
270 let format = format.to_string();
271 tokio::task::spawn_blocking(move || match format.as_str() {
272 "solidworks-part" => Kind::Part(Vec::new()),
273 "solidworks-assembly" => solidworks_assembly(&path).unwrap_or(Kind::Other),
274 "inventor-part" => Kind::Part(Vec::new()),
275 "inventor-assembly" => crate::invplace::kind_of(&path).unwrap_or(Kind::Other),
276 "catia-part" => Kind::Part(Vec::new()),
277 "catia-product" => crate::catprod::kind_of(&path).unwrap_or(Kind::Other),
278 "creo-part" => Kind::Part(Vec::new()),
279 "creo-assembly" => crate::mesh::creo::asm::kind_of(&path).unwrap_or(Kind::Other),
280 _ => nxasm::kind_of(&path).unwrap_or(Kind::Other),
281 })
282 .await
283 .map_err(|e| AppError::Other(anyhow::anyhow!("reading the assembly failed: {e}")))
284}
285
286fn solidworks_assembly(path: &std::path::Path) -> std::io::Result<Kind> {
287 let file = std::io::BufReader::new(std::fs::File::open(path)?);
288 let Some(mut sw) = crate::swpart::SwContainer::open(file)? else {
289 return Ok(Kind::Other);
290 };
291 Ok(
292 match sw
293 .read(crate::swpart::COMPONENT_TREE_SECTION, MAX_TREE)?
294 .and_then(|xml| crate::swpart::placements(&xml))
295 {
296 Some(parts) if !parts.is_empty() => Kind::Assembly(parts),
297 _ => Kind::Other,
298 },
299 )
300}
301
302pub async fn plan(
305 st: &AppState,
306 actor: &Actor,
307 revision_id: Uuid,
308 sha256: &str,
309 format: &str,
310) -> AppResult<Option<Plan>> {
311 let Some(family) = Family::of(format) else {
312 return Ok(None);
313 };
314 let Kind::Assembly(top) = kind(st, sha256, format).await? else {
315 return Ok(None);
316 };
317 let identity = nxasm::identity();
318 let mut plan = Plan {
319 parts: Vec::new(),
320 missing: Vec::new(),
321 key: String::new(),
322 part_format: family.part_format(),
323 };
324 let mut hasher = Sha256::new();
325 hasher.update(sha256.as_bytes());
326 let mut stack = vec![(
329 top,
330 identity,
331 revision_id,
332 vec![sha256.to_string()],
333 Vec::new(),
334 )];
335 while let Some((components, parent, revision, path, overrides)) = stack.pop() {
336 let found = if family == Family::Catia {
337 resolve_catia(st, actor, revision, &components).await?
338 } else {
339 let keys: Vec<String> = components.iter().map(|c| family.key(c)).collect();
340 resolve(&st.db, actor, revision, &keys).await?
341 };
342 let mut kinds: HashMap<String, Kind> = HashMap::new();
343 for c in components.iter().rev() {
344 let placement = Override::pose(&overrides, c).within(&parent);
345 let Some(file) = found.get(&family.key(c)) else {
346 plan.missing.push(c.name.clone());
347 hasher.update(b"\0missing:");
348 hasher.update(c.name.as_bytes());
349 continue;
350 };
351 hasher.update(b"\0");
352 hasher.update(file.sha256.as_bytes());
353 if !family.accepts(&file.format) || path.contains(&file.sha256) {
354 plan.missing.push(c.name.clone());
355 continue;
356 }
357 if !kinds.contains_key(&file.sha256) {
358 let k = kind(st, &file.sha256, &file.format).await?;
359 kinds.insert(file.sha256.clone(), k);
360 }
361 let (own, children) = match kinds.get(&file.sha256) {
363 Some(Kind::Part(children)) => (true, children),
364 Some(Kind::Assembly(children)) => (false, children),
365 _ => {
366 plan.missing.push(c.name.clone());
367 continue;
368 }
369 };
370 if own {
371 if plan.parts.len() >= MAX_PLACED {
372 plan.missing.push(c.name.clone());
373 } else {
374 plan.parts.push((file.sha256.clone(), placement.clone()));
375 }
376 }
377 if !children.is_empty() {
378 if path.len() < MAX_DEPTH {
379 let mut deeper = path.clone();
380 deeper.push(file.sha256.clone());
381 let below = Override::passed_down(&overrides, c);
382 stack.push((children.clone(), placement, file.revision_id, deeper, below));
383 } else {
384 plan.missing.push(c.name.clone());
385 }
386 }
387 }
388 }
389 let digest = hasher.finalize();
390 plan.key = digest.iter().take(12).map(|b| format!("{b:02x}")).collect();
391 Ok(Some(plan))
392}
393
394pub async fn mesh(
397 st: &AppState,
398 actor: &Actor,
399 revision_id: Uuid,
400 sha256: &str,
401 format: &str,
402) -> AppResult<Option<PathBuf>> {
403 let Some(plan) = plan(st, actor, revision_id, sha256, format).await? else {
404 return Ok(None);
405 };
406 let path = st.previews.assembly_path(sha256, &plan.key)?;
407 if tokio::fs::metadata(&path).await.is_ok() {
408 return Ok(Some(path));
409 }
410 let meshes = part_meshes(st, &plan).await?;
411 let target = path.clone();
413 let built = st
414 .previews
415 .render(&path.to_string_lossy(), move || {
416 compose(&plan, &meshes, &target)
417 })
418 .await?
419 .map_err(|e| AppError::Other(anyhow::anyhow!("cannot store the assembly preview: {e}")))?;
420 Ok(built.then_some(path))
421}
422
423async fn part_meshes(st: &AppState, plan: &Plan) -> AppResult<HashMap<String, PathBuf>> {
428 let mut seen = HashSet::new();
429 let distinct: Vec<String> = plan
430 .parts
431 .iter()
432 .filter(|(part, _)| seen.insert(part.as_str()))
433 .map(|(part, _)| part.clone())
434 .collect();
435 let workers = MAX_PARALLEL.min(distinct.len());
436 let queue = Arc::new(Mutex::new(distinct.into_iter()));
437 let mut tasks = tokio::task::JoinSet::new();
439 for _ in 0..workers {
440 let (previews, vault, queue) = (st.previews.clone(), st.vault.clone(), queue.clone());
441 let format = plan.part_format;
442 tasks.spawn(async move {
443 let mut found = Vec::new();
444 loop {
445 let next = queue.lock().unwrap_or_else(|e| e.into_inner()).next();
446 let Some(part) = next else { break };
447 if let Ok(path) = previews.mesh(&vault, &part, format).await {
448 found.push((part, path));
449 }
450 }
451 found
452 });
453 }
454 let mut meshes = HashMap::new();
455 while let Some(found) = tasks.join_next().await {
456 meshes.extend(
457 found.map_err(|e| {
458 AppError::Other(anyhow::anyhow!("converting the parts failed: {e}"))
459 })?,
460 );
461 }
462 Ok(meshes)
463}
464
465fn compose(
469 plan: &Plan,
470 meshes: &HashMap<String, PathBuf>,
471 target: &std::path::Path,
472) -> std::io::Result<bool> {
473 if target.exists() {
474 return Ok(true);
475 }
476 let last: HashMap<&str, usize> = plan
479 .parts
480 .iter()
481 .enumerate()
482 .map(|(at, (part, _))| (part.as_str(), at))
483 .collect();
484 let mut decoded: HashMap<&str, Option<Mesh>> = HashMap::new();
485 let mut out = Mesh::default();
486 for (at, (part, placement)) in plan.parts.iter().enumerate() {
487 let mesh = decoded.entry(part.as_str()).or_insert_with(|| {
488 meshes
489 .get(part)
490 .and_then(|p| std::fs::read(p).ok())
491 .and_then(|b| Mesh::decode(&b).ok())
492 });
493 if let Some(mesh) = mesh.as_ref() {
494 if out.triangle_count() + mesh.triangle_count() > MAX_TRIANGLES {
495 break;
496 }
497 out.append_placed(mesh, |p| placement.place(p));
499 }
500 if last.get(part.as_str()) == Some(&at) {
501 decoded.remove(part.as_str());
502 }
503 }
504 if out.positions.is_empty() || out.check().is_err() {
505 return Ok(false);
506 }
507 let dir = target.parent().unwrap_or(target);
508 std::fs::create_dir_all(dir)?;
509 let tmp = target.with_extension(format!("{}.tmp", Uuid::new_v4().simple()));
510 std::fs::write(&tmp, out.encode())?;
511 std::fs::rename(&tmp, target)?;
512 Ok(true)
513}