Add swizzle variants of with_X methods (#584)

bitshifter · Nov 13, 2024 · c83f129 · c83f129
1 parent f07ed2e
commit c83f129
Show file tree

Hide file tree

Showing 43 changed files with 4,242 additions and 20 deletions.
diff --git a/codegen/templates/swizzle_impl.rs.tera b/codegen/templates/swizzle_impl.rs.tera
@@ -59,6 +59,24 @@ impl Vec{{ dim }}Swizzles for {{ self_t }} {
     #[must_use]
     fn {{ e[j0] }}{{ e[j1] }}(self) -> {{ vec2_t }} {
             {{ vec2_t }} { x: self.{{ e[j0] }}, y: self.{{ e[j1] }} }
+    }
+                    {% endif %}
+                    {% set skip = dim == 2 or j0 == j1 %}
+                    {% if not skip %}
+    #[inline]
+    #[must_use]
+    fn with_{{ e[j0] }}{{ e[j1] }}(self, rhs: {{ vec2_t }}) -> Self {
+        Self::new(
+            {% for l in range(start = 0, end = dim) %}
+                {% if j0 == l %}
+                    rhs.x,
+                {% elif j1 == l %}
+                    rhs.y,
+                {% else %}
+                    self.{{e[l]}},
+                {% endif %}
+            {% endfor %}
+        )
     }
                     {% endif %}
                 {% else %}
@@ -78,6 +96,26 @@ impl Vec{{ dim }}Swizzles for {{ self_t }} {
             {% else %}
                 {{ vec3_t }}::new(self.{{ e[j0] }}, self.{{ e[j1] }}, self.{{ e[j2] }})
             {% endif %}
+    }
+                            {% endif %}
+                            {% set skip = dim == 3 or j0 == j1 or j0 == j2 or j1 == j2 %}
+                            {% if not skip %}
+    #[inline]
+    #[must_use]
+    fn with_{{ e[j0] }}{{ e[j1] }}{{ e[j2] }}(self, rhs: {{ vec3_t }}) -> Self {
+        Self::new(
+            {% for l in range(start = 0, end = dim) %}
+                {% if j0 == l %}
+                    rhs.x,
+                {% elif j1 == l %}
+                    rhs.y,
+                {% elif j2 == l %}
+                    rhs.z,
+                {% else %}
+                    self.{{e[l]}},
+                {% endif %}
+            {% endfor %}
+        )
     }
                             {% endif %}
                         {% else %}

diff --git a/codegen/templates/swizzle_traits.rs.tera b/codegen/templates/swizzle_traits.rs.tera
@@ -3,16 +3,16 @@
 {% set components = ["x", "y", "z", "w"] %}
 {% set dimensions = [2, 3, 4] %}
 {% for dim in dimensions %}
-    {% set ret2 = "Self::Vec2" %}
-    {% set ret3 = "Self::Vec3" %}
-    {% set ret4 = "Self::Vec4" %}
+    {% set vec2_t = "Self::Vec2" %}
+    {% set vec3_t = "Self::Vec3" %}
+    {% set vec4_t = "Self::Vec4" %}
 
     {% if dim == 4 %}
-        {% set ret4 = "Self" %}
+        {% set vec4_t = "Self" %}
     {% elif dim == 3 %}
-        {% set ret3 = "Self" %}
+        {% set vec3_t = "Self" %}
     {% elif dim == 2 %}
-        {% set ret2 = "Self" %}
+        {% set vec2_t = "Self" %}
     {% endif %}
 
     pub trait Vec{{ dim }}Swizzles: Sized + Copy + Clone {
@@ -52,20 +52,28 @@
                     {% if i == 2 %}
                         {% set skip = dim == 2 and e0 == "x" and e1 == "y" %}
                         {% if not skip %}
-                            fn {{e0}}{{e1}}(self) -> {{ret2}};
+                            fn {{e0}}{{e1}}(self) -> {{ vec2_t }};
+                        {% endif %}
+                        {% set skip = dim == 2 or e0 == e1 %}
+                        {% if not skip %}
+                            fn with_{{e0}}{{e1}}(self, rhs: {{ vec2_t }}) -> Self;
                         {% endif %}
                     {% else %}
                         {% for e2 in components | slice(end=dim) %}
                             {% if i == 3 %}
                                 {% set skip = dim == 3 and e0 == "x" and e1 == "y" and e2 == "z" %}
                                 {% if not skip %}
-                                    fn {{e0}}{{e1}}{{e2}}(self) -> {{ret3}};
+                                    fn {{e0}}{{e1}}{{e2}}(self) -> {{ vec3_t }};
+                                {% endif %}
+                                {% set skip = dim == 3 or e0 == e1 or e0 == e2 or e1 == e2 %}
+                                {% if not skip %}
+                                    fn with_{{e0}}{{e1}}{{e2}}(self, rhs: {{ vec3_t }}) -> Self;
                                 {% endif %}
                             {% else %}
                                 {% for e3 in components | slice(end=dim) %}
                                     {% set skip = dim == 4 and e0 == "x" and e1 == "y" and e2 == "z" and e3 == "w" %}
                                     {% if not skip %}
-                                        fn {{e0}}{{e1}}{{e2}}{{e3}}(self) -> {{ret4}};
+                                        fn {{e0}}{{e1}}{{e2}}{{e3}}(self) -> {{ vec4_t }};
                                     {% endif %}
                                 {% endfor %}
                             {% endif %}

diff --git a/src/swizzles/coresimd/vec3a_impl.rs b/src/swizzles/coresimd/vec3a_impl.rs
@@ -29,6 +29,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_xy(self, rhs: Vec2) -> Self {
+        Self::new(rhs.x, rhs.y, self.z)
+    }
+
     #[inline]
     #[must_use]
     fn xz(self) -> Vec2 {
@@ -38,6 +44,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_xz(self, rhs: Vec2) -> Self {
+        Self::new(rhs.x, self.y, rhs.y)
+    }
+
     #[inline]
     #[must_use]
     fn yx(self) -> Vec2 {
@@ -47,6 +59,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_yx(self, rhs: Vec2) -> Self {
+        Self::new(rhs.y, rhs.x, self.z)
+    }
+
     #[inline]
     #[must_use]
     fn yy(self) -> Vec2 {
@@ -65,6 +83,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_yz(self, rhs: Vec2) -> Self {
+        Self::new(self.x, rhs.x, rhs.y)
+    }
+
     #[inline]
     #[must_use]
     fn zx(self) -> Vec2 {
@@ -74,6 +98,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_zx(self, rhs: Vec2) -> Self {
+        Self::new(rhs.y, self.y, rhs.x)
+    }
+
     #[inline]
     #[must_use]
     fn zy(self) -> Vec2 {
@@ -83,6 +113,12 @@ impl Vec3Swizzles for Vec3A {
         }
     }
 
+    #[inline]
+    #[must_use]
+    fn with_zy(self, rhs: Vec2) -> Self {
+        Self::new(self.x, rhs.y, rhs.x)
+    }
+
     #[inline]
     #[must_use]
     fn zz(self) -> Vec2 {