shader: Bug fix: Wrong local variable indices in blocks

internal/shader/shader.go

@@ -633,7 +633,7 @@ func (cs *compileState) parseFunc(block *block, d *ast.FuncDecl) (function, bool
 		}
 	}
 
-	b, ok := cs.parseBlock(block, d.Body.List, inParams, outParams)
+	b, ok := cs.parseBlock(block, d.Name.Name, d.Body.List, inParams, outParams)
 	if !ok {
 		return function{}, false
 	}
@@ -657,17 +657,34 @@ func (cs *compileState) parseFunc(block *block, d *ast.FuncDecl) (function, bool
 	}, true
 }
 
-func (cs *compileState) parseBlock(outer *block, stmts []ast.Stmt, inParams, outParams []variable) (*block, bool) {
+func (cs *compileState) parseBlock(outer *block, fname string, stmts []ast.Stmt, inParams, outParams []variable) (*block, bool) {
 	var vars []variable
 	if outer == &cs.global {
 		vars = make([]variable, 0, len(inParams)+len(outParams))
 		vars = append(vars, inParams...)
 		vars = append(vars, outParams...)
 	}
+
+	var offset int
+	switch {
+	case outer.outer == nil && fname == cs.vertexEntry:
+		offset = 0
+	case outer.outer == nil && fname == cs.fragmentEntry:
+		offset = 0
+	case outer.outer == nil:
+		offset = len(inParams) + len(outParams)
+	case outer.outer.outer == nil:
+		offset = len(outer.outer.vars) + len(outer.vars)
+	default:
+		offset = outer.ir.LocalVarIndexOffset + len(outer.vars)
+	}
+
 	block := &block{
 		vars:  vars,
 		outer: outer,
-		ir:    &shaderir.Block{},
+		ir: &shaderir.Block{
+			LocalVarIndexOffset: offset,
+		},
 	}
 
 	defer func() {
@@ -685,7 +702,7 @@ func (cs *compileState) parseBlock(outer *block, stmts []ast.Stmt, inParams, out
 	}()
 
 	for _, stmt := range stmts {
-		ss, ok := cs.parseStmt(block, stmt, inParams)
+		ss, ok := cs.parseStmt(block, fname, stmt, inParams)
 		if !ok {
 			return nil, false
 		}

internal/shader/stmt.go

@@ -24,7 +24,7 @@ import (
 	"github.com/hajimehoshi/ebiten/internal/shaderir"
 )
 
-func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variable) ([]shaderir.Stmt, bool) {
+func (cs *compileState) parseStmt(block *block, fname string, stmt ast.Stmt, inParams []variable) ([]shaderir.Stmt, bool) {
 	var stmts []shaderir.Stmt
 
 	switch stmt := stmt.(type) {
@@ -97,7 +97,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 			cs.addError(stmt.Pos(), fmt.Sprintf("unexpected token: %s", stmt.Tok))
 		}
 	case *ast.BlockStmt:
-		b, ok := cs.parseBlock(block, stmt.List, inParams, nil)
+		b, ok := cs.parseBlock(block, fname, stmt.List, inParams, nil)
 		if !ok {
 			return nil, false
 		}
@@ -132,7 +132,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		// Create a new pseudo block for the initial statement, so that the counter variable belongs to the
 		// new pseudo block for each for-loop. Without this, the samely named counter variables in different
 		// for-loops confuses the parser.
-		pseudoBlock, ok := cs.parseBlock(block, []ast.Stmt{stmt.Init}, inParams, nil)
+		pseudoBlock, ok := cs.parseBlock(block, fname, []ast.Stmt{stmt.Init}, inParams, nil)
 		if !ok {
 			return nil, false
 		}
@@ -198,7 +198,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		}
 		end := exprs[0].Exprs[1].Const
 
-		postSs, ok := cs.parseStmt(pseudoBlock, stmt.Post, inParams)
+		postSs, ok := cs.parseStmt(pseudoBlock, fname, stmt.Post, inParams)
 		if !ok {
 			return nil, false
 		}
@@ -244,7 +244,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 			return nil, false
 		}
 
-		b, ok := cs.parseBlock(pseudoBlock, []ast.Stmt{stmt.Body}, inParams, nil)
+		b, ok := cs.parseBlock(pseudoBlock, fname, []ast.Stmt{stmt.Body}, inParams, nil)
 		if !ok {
 			return nil, false
 		}
@@ -256,8 +256,9 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		// As the pseudo block is not actually used, copy the variable part to the actual block.
 		// This must be done after parsing the for-loop is done, or the duplicated variables confuses the
 		// parsing.
-		block.vars = append(block.vars, pseudoBlock.vars[0])
+		v := pseudoBlock.vars[0]
-		block.vars[len(block.vars)-1].forLoopCounter = true
+		v.forLoopCounter = true
+		block.vars = append(block.vars, v)
 
 		stmts = append(stmts, shaderir.Stmt{
 			Type:        shaderir.For,
@@ -274,7 +275,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		if stmt.Init != nil {
 			init := stmt.Init
 			stmt.Init = nil
-			b, ok := cs.parseBlock(block, []ast.Stmt{init, stmt}, inParams, nil)
+			b, ok := cs.parseBlock(block, fname, []ast.Stmt{init, stmt}, inParams, nil)
 			if !ok {
 				return nil, false
 			}
@@ -301,7 +302,7 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		stmts = append(stmts, ss...)
 
 		var bs []*shaderir.Block
-		b, ok := cs.parseBlock(block, stmt.Body.List, inParams, nil)
+		b, ok := cs.parseBlock(block, fname, stmt.Body.List, inParams, nil)
 		if !ok {
 			return nil, false
 		}
@@ -310,13 +311,13 @@ func (cs *compileState) parseStmt(block *block, stmt ast.Stmt, inParams []variab
 		if stmt.Else != nil {
 			switch s := stmt.Else.(type) {
 			case *ast.BlockStmt:
-				b, ok := cs.parseBlock(block, s.List, inParams, nil)
+				b, ok := cs.parseBlock(block, fname, s.List, inParams, nil)
 				if !ok {
 					return nil, false
 				}
 				bs = append(bs, b.ir)
 			default:
-				b, ok := cs.parseBlock(block, []ast.Stmt{s}, inParams, nil)
+				b, ok := cs.parseBlock(block, fname, []ast.Stmt{s}, inParams, nil)
 				if !ok {
 					return nil, false
 				}

internal/shader/testdata/blocks3.expected.vs

@@ -0,0 +1,22 @@
+void F0(in vec2 l0, out vec4 l1);
+
+void F0(in vec2 l0, out vec4 l1) {
+	vec4 l2 = vec4(0);
+	{
+		vec3 l3 = vec3(0);
+		vec3 l4 = vec3(0);
+		(l2).x = (l0).x;
+		{
+			vec4 l4 = vec4(0);
+			(l2).y = (l3).y;
+			(l2).z = (l4).z;
+		}
+		{
+			vec4 l5 = vec4(0);
+			(l2).y = (l3).y;
+			(l2).z = (l4).z;
+		}
+	}
+	l1 = l2;
+	return;
+}

internal/shader/testdata/blocks3.go

@@ -0,0 +1,21 @@
+package main
+
+func Foo(foo vec2) vec4 {
+	var r vec4
+	{
+		r.x = foo.x
+		var foo vec3
+		{
+			r.y = foo.y
+			var foo vec4
+			r.z = foo.z
+		}
+		var bar vec3
+		{
+			r.y = foo.y
+			var foo vec4
+			r.z = bar.z
+		}
+	}
+	return r
+}

internal/shader/testdata/for3.expected.vs

@@ -5,7 +5,7 @@ void F0(out vec2 l0) {
 	vec2 l3 = vec2(0);
 	l1 = vec2(0.0);
 	for (int l2 = 0; l2 < 100; l2++) {
-		vec2 l5 = vec2(0);
+		vec2 l3 = vec2(0);
 		l3 = vec2(0.0);
 		l1 = l3;
 	}

internal/shaderir/glsl/glsl.go

@@ -91,7 +91,7 @@ func Compile(p *shaderir.Program) (vertexShader, fragmentShader string) {
 		if p.VertexFunc.Block != nil && len(p.VertexFunc.Block.Stmts) > 0 {
 			vslines = append(vslines, "")
 			vslines = append(vslines, "void main(void) {")
-			vslines = append(vslines, c.glslBlock(p, p.VertexFunc.Block, p.VertexFunc.Block, 0, 0)...)
+			vslines = append(vslines, c.glslBlock(p, p.VertexFunc.Block, p.VertexFunc.Block, 0)...)
 			vslines = append(vslines, "}")
 		}
 	}
@@ -129,7 +129,7 @@ func Compile(p *shaderir.Program) (vertexShader, fragmentShader string) {
 		if p.FragmentFunc.Block != nil && len(p.FragmentFunc.Block.Stmts) > 0 {
 			fslines = append(fslines, "")
 			fslines = append(fslines, "void main(void) {")
-			fslines = append(fslines, c.glslBlock(p, p.FragmentFunc.Block, p.FragmentFunc.Block, 0, 0)...)
+			fslines = append(fslines, c.glslBlock(p, p.FragmentFunc.Block, p.FragmentFunc.Block, 0)...)
 			fslines = append(fslines, "}")
 		}
 	}
@@ -239,7 +239,7 @@ func (c *compileContext) glslFunc(p *shaderir.Program, f *shaderir.Func, prototy
 		return lines
 	}
 	lines = append(lines, fmt.Sprintf("%s {", sig))
-	lines = append(lines, c.glslBlock(p, f.Block, f.Block, 0, idx)...)
+	lines = append(lines, c.glslBlock(p, f.Block, f.Block, 0)...)
 	lines = append(lines, "}")
 
 	return lines
@@ -337,7 +337,7 @@ func localVariable(p *shaderir.Program, topBlock, block *shaderir.Block, idx int
 	}
 }
 
-func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderir.Block, level int, localVarIndex int) []string {
+func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderir.Block, level int) []string {
 	if block == nil {
 		return nil
 	}
@@ -345,8 +345,8 @@ func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderi
 	idt := strings.Repeat("\t", level+1)
 
 	var lines []string
-	for _, t := range block.LocalVars {
+	for i, t := range block.LocalVars {
-		name := fmt.Sprintf("l%d", localVarIndex)
+		name := fmt.Sprintf("l%d", block.LocalVarIndexOffset+i)
 		switch t.Main {
 		case shaderir.Array:
 			lines = append(lines, fmt.Sprintf("%s%s;", idt, c.glslVarDecl(p, &t, name)))
@@ -359,7 +359,6 @@ func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderi
 		default:
 			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, c.glslVarDecl(p, &t, name), c.glslVarInit(p, &t)))
 		}
-		localVarIndex++
 	}
 
 	var glslExpr func(e *shaderir.Expr) string
@@ -420,7 +419,7 @@ func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderi
 			lines = append(lines, fmt.Sprintf("%s%s;", idt, glslExpr(&s.Exprs[0])))
 		case shaderir.BlockStmt:
 			lines = append(lines, idt+"{")
-			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1)...)
 			lines = append(lines, idt+"}")
 		case shaderir.Assign:
 			lhs := s.Exprs[0]
@@ -436,10 +435,10 @@ func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderi
 			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, glslExpr(&lhs), glslExpr(&rhs)))
 		case shaderir.If:
 			lines = append(lines, fmt.Sprintf("%sif (%s) {", idt, glslExpr(&s.Exprs[0])))
-			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1)...)
 			if len(s.Blocks) > 1 {
 				lines = append(lines, fmt.Sprintf("%s} else {", idt))
-				lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[1], level+1, localVarIndex)...)
+				lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[1], level+1)...)
 			}
 			lines = append(lines, fmt.Sprintf("%s}", idt))
 		case shaderir.For:
@@ -482,7 +481,7 @@ func (c *compileContext) glslBlock(p *shaderir.Program, topBlock, block *shaderi
 			end := constantToNumberLiteral(ct, s.ForEnd)
 			t0, t1 := typeString(&t)
 			lines = append(lines, fmt.Sprintf("%sfor (%s %s%s = %s; %s %s %s; %s) {", idt, t0, v, t1, init, v, op, end, delta))
-			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.glslBlock(p, topBlock, s.Blocks[0], level+1)...)
 			lines = append(lines, fmt.Sprintf("%s}", idt))
 		case shaderir.Continue:
 			lines = append(lines, idt+"continue;")

internal/shaderir/ir_test.go

@@ -23,10 +23,11 @@ import (
 	"github.com/hajimehoshi/ebiten/internal/shaderir/metal"
 )
 
-func block(localVars []Type, stmts ...Stmt) *Block {
+func block(localVars []Type, offset int, stmts ...Stmt) *Block {
 	return &Block{
-		LocalVars: localVars,
+		LocalVars:           localVars,
-		Stmts:     stmts,
+		LocalVarIndexOffset: offset,
+		Stmts:               stmts,
 	}
 }
 
@@ -279,6 +280,7 @@ void F0(in float l0, in vec2 l1, in vec4 l2, out mat4 l3) {
 						Return: Type{Main: Float},
 						Block: block(
 							nil,
+							0,
 							returnStmt(
 								localVariableExpr(0),
 							),
@@ -313,7 +315,7 @@ float F0(in float l0) {
 						Block: block([]Type{
 							{Main: Mat4},
 							{Main: Mat4},
-						}),
+						}, 2),
 					},
 				},
 			},
@@ -348,12 +350,14 @@ void F0(in float l0, out float l1) {
 								{Main: Mat4},
 								{Main: Mat4},
 							},
+							2,
 							blockStmt(
 								block(
 									[]Type{
 										{Main: Mat4},
 										{Main: Mat4},
 									},
+									4,
 								),
 							),
 						),
@@ -397,6 +401,7 @@ void F0(in float l0, out float l1) {
 						},
 						Block: block(
 							nil,
+							0,
 							assignStmt(
 								localVariableExpr(2),
 								binaryExpr(
@@ -437,6 +442,7 @@ void F0(in float l0, in float l1, out float l2) {
 						},
 						Block: block(
 							nil,
+							0,
 							assignStmt(
 								localVariableExpr(3),
 								selectionExpr(
@@ -476,6 +482,7 @@ void F0(in bool l0, in float l1, in float l2, out float l3) {
 						},
 						Block: block(
 							nil,
+							0,
 							exprStmt(
 								callExpr(
 									functionExpr(1),
@@ -522,6 +529,7 @@ void F0(in float l0, in float l1, out vec2 l2) {
 						},
 						Block: block(
 							nil,
+							0,
 							assignStmt(
 								localVariableExpr(2),
 								callExpr(
@@ -560,6 +568,7 @@ void F0(in float l0, in float l1, out float l2) {
 						},
 						Block: block(
 							nil,
+							0,
 							assignStmt(
 								localVariableExpr(1),
 								fieldSelectorExpr(
@@ -598,6 +607,7 @@ void F0(in vec4 l0, out vec2 l1) {
 						},
 						Block: block(
 							nil,
+							0,
 							ifStmt(
 								binaryExpr(
 									EqualOp,
@@ -606,6 +616,7 @@ void F0(in vec4 l0, out vec2 l1) {
 								),
 								block(
 									nil,
+									0,
 									assignStmt(
 										localVariableExpr(2),
 										localVariableExpr(0),
@@ -613,6 +624,7 @@ void F0(in vec4 l0, out vec2 l1) {
 								),
 								block(
 									nil,
+									0,
 									assignStmt(
 										localVariableExpr(2),
 										localVariableExpr(1),
@@ -660,6 +672,7 @@ void F0(in float l0, in float l1, out float l2) {
 							[]Type{
 								{},
 							},
+							0,
 							forStmt(
 								Type{Main: Int},
 								3,
@@ -669,6 +682,7 @@ void F0(in float l0, in float l1, out float l2) {
 								1,
 								block(
 									nil,
+									4,
 									assignStmt(
 										localVariableExpr(2),
 										localVariableExpr(0),
@@ -712,6 +726,7 @@ void F0(in float l0, in float l1, out float l2) {
 							[]Type{
 								{},
 							},
+							0,
 							forStmt(
 								Type{Main: Int},
 								3,
@@ -723,6 +738,7 @@ void F0(in float l0, in float l1, out float l2) {
 									[]Type{
 										{Main: Int},
 									},
+									4,
 									assignStmt(
 										localVariableExpr(2),
 										localVariableExpr(4),
@@ -783,6 +799,7 @@ void F0(float l0, float l1, thread float& l2) {
 								{},
 								{},
 							},
+							0,
 							forStmt(
 								Type{Main: Int},
 								3,
@@ -794,9 +811,10 @@ void F0(float l0, float l1, thread float& l2) {
 									[]Type{
 										{Main: Int},
 									},
+									4,
 									assignStmt(
 										localVariableExpr(2),
-										localVariableExpr(5),
+										localVariableExpr(4),
 									),
 								),
 							),
@@ -811,6 +829,7 @@ void F0(float l0, float l1, thread float& l2) {
 									[]Type{
 										{Main: Int},
 									},
+									5,
 									assignStmt(
 										localVariableExpr(2),
 										localVariableExpr(5),
@@ -825,8 +844,8 @@ void F0(float l0, float l1, thread float& l2) {
 
 void F0(in float l0, in float l1, out float l2) {
 	for (int l3 = 0; l3 < 100; l3++) {
-		int l5 = 0;
+		int l4 = 0;
-		l2 = l5;
+		l2 = l4;
 	}
 	for (float l4 = 0.0; l4 < 100.0; l4++) {
 		int l5 = 0;
@@ -838,8 +857,8 @@ void F0(in float l0, in float l1, out float l2);
 
 void F0(in float l0, in float l1, out float l2) {
 	for (int l3 = 0; l3 < 100; l3++) {
-		int l5 = 0;
+		int l4 = 0;
-		l2 = l5;
+		l2 = l4;
 	}
 	for (float l4 = 0.0; l4 < 100.0; l4++) {
 		int l5 = 0;
@@ -856,8 +875,8 @@ void F0(float l0, float l1, thread float& l2);
 
 void F0(float l0, float l1, thread float& l2) {
 	for (int l3 = 0; l3 < 100; l3++) {
-		int l5 = 0;
+		int l4 = 0;
-		l2 = l5;
+		l2 = l4;
 	}
 	for (float l4 = 0.0; l4 < 100.0; l4++) {
 		int l5 = 0;
@@ -883,6 +902,7 @@ void F0(float l0, float l1, thread float& l2) {
 				VertexFunc: VertexFunc{
 					Block: block(
 						nil,
+						0,
 						assignStmt(
 							localVariableExpr(3),
 							localVariableExpr(0),
@@ -933,6 +953,7 @@ varying vec2 V1;`,
 				VertexFunc: VertexFunc{
 					Block: block(
 						nil,
+						0,
 						assignStmt(
 							localVariableExpr(3),
 							localVariableExpr(0),
@@ -953,6 +974,7 @@ varying vec2 V1;`,
 							{Main: Float},
 							{Main: Vec2},
 						},
+						0,
 						assignStmt(
 							localVariableExpr(3),
 							localVariableExpr(0),

internal/shaderir/metal/metal.go

@@ -111,7 +111,7 @@ func Compile(p *shaderir.Program, vertex, fragment string) (shader string) {
 		}
 		lines[len(lines)-1] += ") {"
 		lines = append(lines, fmt.Sprintf("\tVaryings %s = {};", vertexOut))
-		lines = append(lines, c.metalBlock(p, p.VertexFunc.Block, p.VertexFunc.Block, 0, 0)...)
+		lines = append(lines, c.metalBlock(p, p.VertexFunc.Block, p.VertexFunc.Block, 0)...)
 		if last := fmt.Sprintf("\treturn %s;", vertexOut); lines[len(lines)-1] != last {
 			lines = append(lines, last)
 		}
@@ -133,7 +133,7 @@ func Compile(p *shaderir.Program, vertex, fragment string) (shader string) {
 		}
 		lines[len(lines)-1] += ") {"
 		lines = append(lines, fmt.Sprintf("\tfloat4 %s = float4(0);", fragmentOut))
-		lines = append(lines, c.metalBlock(p, p.FragmentFunc.Block, p.FragmentFunc.Block, 0, 0)...)
+		lines = append(lines, c.metalBlock(p, p.FragmentFunc.Block, p.FragmentFunc.Block, 0)...)
 		if last := fmt.Sprintf("\treturn %s;", fragmentOut); lines[len(lines)-1] != last {
 			lines = append(lines, last)
 		}
@@ -245,7 +245,7 @@ func (c *compileContext) metalFunc(p *shaderir.Program, f *shaderir.Func, protot
 		return lines
 	}
 	lines = append(lines, fmt.Sprintf("%s {", sig))
-	lines = append(lines, c.metalBlock(p, f.Block, f.Block, 0, idx)...)
+	lines = append(lines, c.metalBlock(p, f.Block, f.Block, 0)...)
 	lines = append(lines, "}")
 
 	return lines
@@ -311,7 +311,7 @@ func localVariableName(p *shaderir.Program, topBlock *shaderir.Block, idx int) s
 	}
 }
 
-func (c *compileContext) metalBlock(p *shaderir.Program, topBlock, block *shaderir.Block, level int, localVarIndex int) []string {
+func (c *compileContext) metalBlock(p *shaderir.Program, topBlock, block *shaderir.Block, level int) []string {
 	if block == nil {
 		return nil
 	}
@@ -319,12 +319,11 @@ func (c *compileContext) metalBlock(p *shaderir.Program, topBlock, block *shader
 	idt := strings.Repeat("\t", level+1)
 
 	var lines []string
-	for _, t := range block.LocalVars {
+	for i, t := range block.LocalVars {
 		// The type is None e.g., when the variable is a for-loop counter.
 		if t.Main != shaderir.None {
-			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, c.metalVarDecl(p, &t, fmt.Sprintf("l%d", localVarIndex), false, false), c.metalVarInit(p, &t)))
+			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, c.metalVarDecl(p, &t, fmt.Sprintf("l%d", block.LocalVarIndexOffset+i), false, false), c.metalVarInit(p, &t)))
 		}
-		localVarIndex++
 	}
 
 	var metalExpr func(e *shaderir.Expr) string
@@ -395,16 +394,16 @@ func (c *compileContext) metalBlock(p *shaderir.Program, topBlock, block *shader
 			lines = append(lines, fmt.Sprintf("%s%s;", idt, metalExpr(&s.Exprs[0])))
 		case shaderir.BlockStmt:
 			lines = append(lines, idt+"{")
-			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1)...)
 			lines = append(lines, idt+"}")
 		case shaderir.Assign:
 			lines = append(lines, fmt.Sprintf("%s%s = %s;", idt, metalExpr(&s.Exprs[0]), metalExpr(&s.Exprs[1])))
 		case shaderir.If:
 			lines = append(lines, fmt.Sprintf("%sif (%s) {", idt, metalExpr(&s.Exprs[0])))
-			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1)...)
 			if len(s.Blocks) > 1 {
 				lines = append(lines, fmt.Sprintf("%s} else {", idt))
-				lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[1], level+1, localVarIndex)...)
+				lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[1], level+1)...)
 			}
 			lines = append(lines, fmt.Sprintf("%s}", idt))
 		case shaderir.For:
@@ -447,7 +446,7 @@ func (c *compileContext) metalBlock(p *shaderir.Program, topBlock, block *shader
 			end := constantToNumberLiteral(ct, s.ForEnd)
 			ts := typeString(&t, false, false)
 			lines = append(lines, fmt.Sprintf("%sfor (%s %s = %s; %s %s %s; %s) {", idt, ts, v, init, v, op, end, delta))
-			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1, localVarIndex)...)
+			lines = append(lines, c.metalBlock(p, topBlock, s.Blocks[0], level+1)...)
 			lines = append(lines, fmt.Sprintf("%s}", idt))
 		case shaderir.Continue:
 			lines = append(lines, idt+"continue;")

internal/shaderir/program.go

@@ -57,8 +57,9 @@ type FragmentFunc struct {
 }
 
 type Block struct {
-	LocalVars []Type
+	LocalVars           []Type
-	Stmts     []Stmt
+	LocalVarIndexOffset int
+	Stmts               []Stmt
 }
 
 type Stmt struct {

internal/testing/shader.go

@@ -174,6 +174,7 @@ var (
 	}
 	defaultVertexFunc = shaderir.VertexFunc{
 		Block: &shaderir.Block{
+			LocalVarIndexOffset: 0,
 			Stmts: []shaderir.Stmt{
 				{
 					Type: shaderir.Assign,
@@ -280,6 +281,7 @@ func ShaderProgramFill(r, g, b, a byte) shaderir.Program {
 	p := defaultProgram()
 	p.FragmentFunc = shaderir.FragmentFunc{
 		Block: &shaderir.Block{
+			LocalVarIndexOffset: 0,
 			Stmts: []shaderir.Stmt{
 				{
 					Type: shaderir.Assign,
@@ -417,7 +419,8 @@ func ShaderProgramImages(imageNum int) shaderir.Program {
 			LocalVars: []shaderir.Type{
 				{Main: shaderir.Vec4},
 			},
-			Stmts: stmts,
+			LocalVarIndexOffset: 0,
+			Stmts:               stmts,
 		},
 	}