Skip to main content

teamcraft/services/
assemblies.rs

1//! 3D previews of NX, SOLIDWORKS, Inventor and Creo assemblies and CATIA
2//! products.
3//!
4//! An assembly's preview is put together from the previews of its
5//! component files, with the placements the assembly stores: NX in its
6//! structure stream (`nxasm`), SOLIDWORKS in its component tree
7//! (`swpart::placements`, sub-assemblies already resolved there), Inventor
8//! in `UFRxDoc` and `AmRxSegment` (`invplace`, sub-assemblies are files of
9//! their own; a parent may re-pose children of a flexible
10//! sub-assembly, `invplace::Override`), CATIA in the object stream of the
11//! product (`catprod`, sub-products are files of their own), Creo in the
12//! assembly's `NeuAsmSld` section (`mesh::creo::asm`, sub-assemblies are
13//! files of their own).
14//!
15//! 1. The components are found **by file name** (`NAME.prt` for NX, the
16//!    file name of the stored path for SOLIDWORKS and Inventor, `name.prt`
17//!    or `name.asm` for Creo, whose files are stored without their
18//!    iteration suffix), the way file
19//!    references are matched (`references::files_of_revision`): a file of
20//!    the assembly's own revision first, otherwise the latest revision of
21//!    another item the viewer may see. A CATIA product names no files: a
22//!    component is the document whose `CATProdCont` UUID it refers to
23//!    (among the CATIA files of the product's revision and those found by
24//!    name), otherwise `<part number>.CATPart` or `.CATProduct` by name.
25//!    Components that are not in TeamCraft are left out.
26//! 2. A component that is an assembly itself is resolved the same way,
27//!    its placements composed with the parent's (depth and cycles bounded).
28//! 3. Every part is converted like any other part (isolated preview
29//!    process, cache per content; as many at a time as the preview queue
30//!    runs), then every instance is placed straight into the assembly mesh.
31//!    The composition runs through `Previews::render`: once, however many
32//!    viewers and thumbnails ask for the assembly at the same time.
33//!
34//! The result depends on which component versions were found, so it is
35//! cached under a key built from the content hashes of everything used
36//! (`Previews::assembly_path`): a new version of a component gives a new
37//! preview, an unchanged assembly is built once.
38//!
39//! # Invariants
40//!
41//! * **Permission:** only files the viewer may see are used (the same
42//!   project-group filter as file references); nothing is written to the
43//!   database.
44//! * **Bounds:** at most [`MAX_DEPTH`] levels, [`MAX_PLACED`] placed parts
45//!   and `mesh::MAX_TRIANGLES` triangles; an assembly that contains itself
46//!   is not entered again.
47
48use 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
69/// Nesting levels of sub-assemblies followed at most.
70pub const MAX_DEPTH: usize = 8;
71/// Placed parts at most.
72pub const MAX_PLACED: usize = 50_000;
73
74/// A component file found in the vault.
75#[derive(Clone, Debug, sqlx::FromRow)]
76struct Found {
77    key: String,
78    sha256: String,
79    format: String,
80    revision_id: Uuid,
81}
82
83/// The CAD system of an assembly: how its components are named and which
84/// files may be components.
85#[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    /// The lower-case file name of a component (CATIA: the UUID of its
107    /// document, else its part number).
108    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    /// The preview format of the parts.
131    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/// The parts of an assembly with their placements in the assembly.
143#[derive(Debug)]
144pub struct Plan {
145    /// (content hash of the part file, placement in the top assembly).
146    pub parts: Vec<(String, nxasm::Component)>,
147    /// Component names that were not found (or could not be read).
148    pub missing: Vec<String>,
149    /// Hex key of everything used (for the cache).
150    pub key: String,
151    /// The preview format of the parts.
152    pub part_format: PreviewFormat,
153}
154
155/// Finds the component files by their lower-case file names.
156async 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        // the first match wins: own revision, then the lowest item number
189        out.entry(row.key.clone()).or_insert(row);
190    }
191    Ok(out)
192}
193
194/// CATIA files of a product's own revision whose documents are compared.
195const MAX_OWN_CATIA: i64 = 2_000;
196
197/// Finds the documents of a CATIA product's components, keyed like
198/// [`Family::key`]: by the UUID of their `CATProdCont` among the CATIA
199/// files of the product's revision and the files found by name, otherwise
200/// `<part number>.CATPart` / `.CATProduct` by name. Copies of a document
201/// keep its UUID: a file of the product's own revision wins.
202async 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    // the documents, read off the first bytes of every candidate
232    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
261/// Largest component tree read from a SOLIDWORKS assembly.
262const MAX_TREE: u64 = 64 * 1024 * 1024;
263
264/// What the file of a blob is (read with seeks, off the async runtime). A
265/// SOLIDWORKS or Inventor part has no components; a SOLIDWORKS assembly
266/// lists its parts, sub-assemblies resolved; an Inventor assembly lists its
267/// visible parts and sub-assemblies.
268async 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
302/// The parts of the assembly `sha256` (a file of `revision_id` in the
303/// format `format`), or `None` when it is no assembly.
304pub 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    // (components, placement of the assembly, its revision, path of hashes,
327    // poses the levels above give components further down – Inventor)
328    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            // a part may hold components besides its own bodies
362            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
394/// The composed preview mesh of an assembly (cached), or `None` when the
395/// file is no assembly or none of its parts is in TeamCraft.
396pub 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    // composed once, however many viewers (and the thumbnail) ask for it
412    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
423/// The preview mesh of every distinct part, through the normal (isolated,
424/// cached) preview: as many conversions at a time as the preview queue runs
425/// ([`MAX_PARALLEL`]) instead of one after the other. A part that has no
426/// preview is left out.
427async 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    // the tasks end with this request: a dropped `JoinSet` stops them
438    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
465/// Places the part meshes and writes the assembly mesh to `target`;
466/// `false` when no part could be placed. Run through `Previews::render`:
467/// a request that waited for another one's composition finds `target`.
468fn 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    // where every part is placed for the last time: its mesh is given up
477    // there, so only the parts still to be placed are held in memory
478    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            // placed straight into the assembly, without a copy of the part
498            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}