internal/shader: bug fix: wrong local variable index was used for assignments

Closes #2655

internal/shader/stmt.go

@@ -558,7 +558,7 @@ func (cs *compileState) assign(block *block, fname string, pos token.Pos, lhs, r
 				block.vars = append(block.vars, variable{
 					typ: t,
 				})
-				idx := len(block.vars) - 1
+				idx := block.totalLocalVariableCount() - 1
 				stmts = append(stmts,
 					shaderir.Stmt{
 						Type: shaderir.Assign,

internal/shader/testdata/issue2655.expected.vs

@@ -0,0 +1,43 @@
+float F0(void);
+void F1(out float l0, out vec2 l1);
+void F2(in bool l0, out float l1, out vec2 l2);
+
+float F0(void) {
+	float l0 = float(0);
+	{
+		l0 = 0;
+	}
+	return l0;
+}
+
+void F1(out float l0, out vec2 l1) {
+	float l2 = float(0);
+	vec2 l3 = vec2(0);
+	{
+		int l4 = 0;
+		vec2 l5 = vec2(0);
+		l4 = 0;
+		l2 = l4;
+		l5 = vec2(0.0);
+		l3 = l5;
+	}
+	l0 = l2;
+	l1 = l3;
+	return;
+}
+
+void F2(in bool l0, out float l1, out vec2 l2) {
+	float l3 = float(0);
+	vec2 l4 = vec2(0);
+	{
+		int l5 = 0;
+		vec2 l6 = vec2(0);
+		l5 = 0;
+		l3 = l5;
+		l6 = vec2(0.0);
+		l4 = l6;
+	}
+	l1 = l3;
+	l2 = l4;
+	return;
+}

internal/shader/testdata/issue2655.go

@@ -0,0 +1,27 @@
+package main
+
+func Ret1() float {
+	var a float
+	{
+		a = 0
+	}
+	return a
+}
+
+func Ret2() (float, vec2) {
+	var a float
+	var b vec2
+	{
+		a, b = 0, vec2(0)
+	}
+	return a, b
+}
+
+func Ret2Plus(foo bool) (float, vec2) {
+	var a float
+	var b vec2
+	{
+		a, b = 0, vec2(0)
+	}
+	return a, b
+}