From 4a8539112c5473a6e6202b08e9958bd3dcbe78cf Mon Sep 17 00:00:00 2001
From: Kayh <basedlainpilled@gmail.com>
Date: Tue, 12 Mar 2024 20:57:29 -0400
Subject: [PATCH] add bevy primitive import extension

---
 bevy_gltf_kun/src/import/extensions.rs     | 52 ++++++++++++++++------
 bevy_gltf_kun/src/import/gltf/mesh.rs      |  8 ++--
 bevy_gltf_kun/src/import/gltf/node.rs      |  3 +-
 bevy_gltf_kun/src/import/gltf/primitive.rs |  7 ++-
 4 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/bevy_gltf_kun/src/import/extensions.rs b/bevy_gltf_kun/src/import/extensions.rs
index 7bcd27c..65435d0 100644
--- a/bevy_gltf_kun/src/import/extensions.rs
+++ b/bevy_gltf_kun/src/import/extensions.rs
@@ -2,23 +2,13 @@ use bevy::prelude::*;
 use gltf_kun::{
     extensions::{DefaultExtensions, Extension},
     graph::{
-        gltf::{document::GltfDocument, node::Node},
+        gltf::{GltfDocument, Node, Primitive},
         Extensions,
     },
 };
 
 use crate::import::gltf::document::ImportContext;
 
-pub trait RootExtensionImport<D: Extensions>: Extension {
-    fn maybe_import_root(context: &mut ImportContext) {
-        if let Some(ext) = context.doc.get_extension::<Self>(context.graph) {
-            Self::import_root(context, ext);
-        }
-    }
-
-    fn import_root(context: &mut ImportContext, ext: Self);
-}
-
 pub trait NodeExtensionImport<D>: Extension {
     fn maybe_import_node(context: &mut ImportContext, entity: &mut EntityWorldMut, node: Node) {
         if let Some(ext) = node.get_extension::<Self>(context.graph) {
@@ -31,13 +21,47 @@ pub trait NodeExtensionImport<D>: Extension {
     fn import_node(context: &mut ImportContext, entity: &mut EntityWorldMut, ext: Self);
 }
 
+pub trait PrimitiveExtensionImport<D>: Extension {
+    fn maybe_import_primitive(
+        context: &mut ImportContext,
+        entity: &mut EntityWorldMut,
+        primitive: Primitive,
+    ) {
+        if let Some(ext) = primitive.get_extension::<Self>(context.graph) {
+            Self::import_primitive(context, entity, ext);
+        }
+    }
+
+    fn import_primitive(context: &mut ImportContext, entity: &mut EntityWorldMut, ext: Self);
+}
+
+pub trait RootExtensionImport<D: Extensions>: Extension {
+    fn maybe_import_root(context: &mut ImportContext) {
+        if let Some(ext) = context.doc.get_extension::<Self>(context.graph) {
+            Self::import_root(context, ext);
+        }
+    }
+
+    fn import_root(context: &mut ImportContext, ext: Self);
+}
+
 pub trait BevyImportExtensions<D> {
-    fn import_root(context: &mut ImportContext);
     fn import_node(context: &mut ImportContext, entity: &mut EntityWorldMut, node: Node);
+    fn import_primitive(
+        context: &mut ImportContext,
+        entity: &mut EntityWorldMut,
+        primitive: Primitive,
+    );
+    fn import_root(context: &mut ImportContext);
 }
 
 impl BevyImportExtensions<GltfDocument> for DefaultExtensions {
-    fn import_root(_context: &mut ImportContext) {}
+    fn import_primitive(
+        _context: &mut ImportContext,
+        _entity: &mut EntityWorldMut,
+        _primitive: Primitive,
+    ) {
+    }
 
     fn import_node(context: &mut ImportContext, entity: &mut EntityWorldMut, node: Node) {
         #[cfg(feature = "omi_physics")]
@@ -47,4 +71,6 @@ impl BevyImportExtensions<GltfDocument> for DefaultExtensions {
             )
         }
     }
+
+    fn import_root(_context: &mut ImportContext) {}
 }
diff --git a/bevy_gltf_kun/src/import/gltf/mesh.rs b/bevy_gltf_kun/src/import/gltf/mesh.rs
index a446a28..2f69f1e 100644
--- a/bevy_gltf_kun/src/import/gltf/mesh.rs
+++ b/bevy_gltf_kun/src/import/gltf/mesh.rs
@@ -1,9 +1,11 @@
 use bevy::prelude::*;
 use gltf_kun::graph::{
-    gltf::{self},
+    gltf::{self, GltfDocument},
     GraphNodeWeight,
 };
 
+use crate::import::extensions::BevyImportExtensions;
+
 use super::{
     document::ImportContext,
     primitive::{import_primitive, GltfPrimitive},
@@ -15,7 +17,7 @@ pub struct GltfMesh {
     pub extras: Option<Box<serde_json::value::RawValue>>,
 }
 
-pub fn import_mesh(
+pub fn import_mesh<E: BevyImportExtensions<GltfDocument>>(
     context: &mut ImportContext,
     entity: &mut EntityWorldMut,
     mut m: gltf::mesh::Mesh,
@@ -31,7 +33,7 @@ pub fn import_mesh(
 
     entity.with_children(|parent| {
         for (i, p) in m.primitives(context.graph).iter_mut().enumerate() {
-            match import_primitive(context, parent, is_scale_inverted, m, &mesh_label, i, p) {
+            match import_primitive::<E>(context, parent, is_scale_inverted, m, &mesh_label, i, p) {
                 Ok((ent, handle, weights)) => {
                     morph_weights = weights;
                     primitive_entities.push(ent);
diff --git a/bevy_gltf_kun/src/import/gltf/node.rs b/bevy_gltf_kun/src/import/gltf/node.rs
index c010d28..ccfdcaa 100644
--- a/bevy_gltf_kun/src/import/gltf/node.rs
+++ b/bevy_gltf_kun/src/import/gltf/node.rs
@@ -55,7 +55,8 @@ pub fn import_node<E: BevyImportExtensions<GltfDocument>>(
 
     let mesh = match n.mesh(context.graph) {
         Some(m) => {
-            let (ents, mesh, morph_weights) = import_mesh(context, &mut ent, m, is_scale_inverted);
+            let (ents, mesh, morph_weights) =
+                import_mesh::<E>(context, &mut ent, m, is_scale_inverted);
 
             primitive_entities.extend(ents);
 
diff --git a/bevy_gltf_kun/src/import/gltf/primitive.rs b/bevy_gltf_kun/src/import/gltf/primitive.rs
index bd9d88a..6104ad5 100644
--- a/bevy_gltf_kun/src/import/gltf/primitive.rs
+++ b/bevy_gltf_kun/src/import/gltf/primitive.rs
@@ -22,11 +22,14 @@ use gltf_kun::graph::{
             Accessor, ComponentType, GetAccessorSliceError, Type,
         },
         primitive::{Mode, MorphTarget, Primitive, Semantic},
+        GltfDocument,
     },
     Graph, GraphNodeWeight,
 };
 use thiserror::Error;
 
+use crate::import::extensions::BevyImportExtensions;
+
 use super::{
     document::ImportContext,
     material::{default_material, import_material},
@@ -61,7 +64,7 @@ pub enum ImportPrimitiveError {
     MorphBuildError(#[from] MorphBuildError),
 }
 
-pub fn import_primitive(
+pub fn import_primitive<E: BevyImportExtensions<GltfDocument>>(
     context: &mut ImportContext,
     parent: &mut WorldChildBuilder,
     is_scale_inverted: bool,
@@ -211,6 +214,8 @@ pub fn import_primitive(
         bevy_mesh.set_morph_targets(handle);
     }
 
+    E::import_primitive(context, &mut entity, *p);
+
     let mesh_label = context
         .load_context
         .add_labeled_asset(primitive_label.clone(), bevy_mesh);