From f4532411833f3f7f79d721c569196ed1a0abc823 Mon Sep 17 00:00:00 2001
From: Josh Watzman <github@jwatzman.org>
Date: Sun, 16 Apr 2023 21:32:52 +0100
Subject: [PATCH] Fix potential unsafe swizzle optimizations

The optimization in #253 is unsafe -- `vec3(a.x)` and `vec3(x)` are not
equivalent. For a single-element swizzles, they are only sure to be safe
if they are the Nth argument to a vecN constructor. This change counts
the number of arguments, and only allows a single-element swizzle to be
optimized in that case.

This does miss some edge cases -- for example, I think the swizzle in
code like `vec3(foo.xy, bar.x)` is safe to remove (`vec3(foo.xy, bar)`)
but that doesn't seem to actually happen in any of the examples in this
repo, and it's super tricky to get right in practise (if `foo.xy` is
some variable instead of a swizzle right there), etc etc. So there is
potentially still something on the table, but at least for the cases
available this is good enough for now.

Fixes #297
---
 src/rewriter.fs                  |  9 +++++----
 tests/unit/vectors.frag          | 11 ++++++++++-
 tests/unit/vectors.frag.expected | 10 +++++++++-
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/rewriter.fs b/src/rewriter.fs
index c69ecbd2..9ee601aa 100644
--- a/src/rewriter.fs
+++ b/src/rewriter.fs
@@ -247,17 +247,18 @@ let private simplifyVec (constr: Ident) args =
         | _ -> allArgs
 
     // vec3(a.x, b.xy) => vec3(a.x, b)
-    let rec dropLastSwizzle = function
+    let rec dropLastSwizzle n = function
         | [Dot (expr, field) as last] when isFieldSwizzle field ->
             match [for c in field -> swizzleIndex c] with
-            | [0] | [0; 1] | [0; 1; 2] | [0; 1; 2; 3] -> [expr]
+            | [0] when n = 1 -> [expr]
+            | [0; 1] | [0; 1; 2] | [0; 1; 2; 3] -> [expr]
             | _ -> [last]
-        | e1 :: rest -> e1 :: dropLastSwizzle rest
+        | e1 :: rest -> e1 :: dropLastSwizzle (n-1) rest
         | x -> x
 
     let args = combineSwizzles args |> List.map useInts
     let args = if args.Length = vecSize then mergeAllEquals args args else args
-    let args = dropLastSwizzle args
+    let args = dropLastSwizzle vecSize args
     FunCall (Var constr, args)
 
 let private simplifyExpr (didInline: bool ref) env = function
diff --git a/tests/unit/vectors.frag b/tests/unit/vectors.frag
index a2c4902b..2cbbf866 100644
--- a/tests/unit/vectors.frag
+++ b/tests/unit/vectors.frag
@@ -23,7 +23,16 @@ float withExtraComponents() {
     vec2 v2 = vec2(1);
     vec3 v3 = vec3(1, 2, v2.x);
     vec4 v4 = vec4(v2, v3.xy);
-    return v2.x + v3.x + v4.x;
+    vec3 v5 = vec3(1, v3.rg);
+    return v2.x + v3.x + v4.x + v5.x;
+}
+
+float withExtraComponentsUnsafe(in vec3 a) {
+    vec3 v1 = vec3(a.x);
+    vec3 v2 = vec3(1.0, a.x);
+    vec2 v3 = vec3(1.0, a.y);
+    vec3 v4 = vec3(1.0, 2.0, a.y);
+    return v1.x + v2.x + v3.x + v4.x;
 }
 
 struct S{
diff --git a/tests/unit/vectors.frag.expected b/tests/unit/vectors.frag.expected
index a6bc2ab1..eafec719 100644
--- a/tests/unit/vectors.frag.expected
+++ b/tests/unit/vectors.frag.expected
@@ -21,7 +21,15 @@ float withExtraComponents()
   vec2 v2=vec2(1);
   vec3 v3=vec3(1,2,v2);
   vec4 v4=vec4(v2,v3);
-  return v2.x+v3.x+v4.x;
+  vec3 v5=vec3(1,v3);
+  return v2.x+v3.x+v4.x+v5.x;
+}
+float withExtraComponentsUnsafe(vec3 a)
+{
+  vec3 v1=vec3(a.x),v2=vec3(1,a.x);
+  vec2 v3=vec3(1,a.y);
+  vec3 v4=vec3(1,2,a.y);
+  return v1.x+v2.x+v3.x+v4.x;
 }struct S{vec2 p1;vec2 cp1;vec2 cp2;vec2 p2;};
 vec2 calc(S points,float t)
 {