From 4a8bd688a98ae466c23473fde5a1289f5f037879 Mon Sep 17 00:00:00 2001
From: Hajime Hoshi <hajimehoshi@gmail.com>
Date: Sun, 12 Jul 2020 18:40:01 +0900
Subject: [PATCH] shader: Bug fix: index out of range at defining a var

---
 internal/shader/shader.go                    | 15 +++++++++++++++
 internal/shader/stmt.go                      |  7 ++++++-
 internal/shader/testdata/blocks2.expected.vs | 12 ++++++++++++
 internal/shader/testdata/blocks2.go          | 10 ++++++++++
 4 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 internal/shader/testdata/blocks2.expected.vs
 create mode 100644 internal/shader/testdata/blocks2.go

diff --git a/internal/shader/shader.go b/internal/shader/shader.go
index c1f6c32ea..ee5fd37dc 100644
--- a/internal/shader/shader.go
+++ b/internal/shader/shader.go
@@ -110,6 +110,21 @@ func (b *block) findLocalVariable(name string) (int, shaderir.Type, bool) {
 	return 0, shaderir.Type{}, false
 }
 
+func (b *block) findLocalVariableByIndex(idx int) (shaderir.Type, bool) {
+	bs := []*block{b}
+	for outer := b.outer; outer != nil; outer = outer.outer {
+		bs = append(bs, outer)
+	}
+	for i := len(bs) - 1; i >= 0; i-- {
+		if len(bs[i].vars) <= idx {
+			idx -= len(bs[i].vars)
+			continue
+		}
+		return bs[i].vars[idx].typ, true
+	}
+	return shaderir.Type{}, false
+}
+
 type ParseError struct {
 	errs []string
 }
diff --git a/internal/shader/stmt.go b/internal/shader/stmt.go
index 37bd2a13e..808ba658e 100644
--- a/internal/shader/stmt.go
+++ b/internal/shader/stmt.go
@@ -268,7 +268,12 @@ func (cs *compileState) assign(block *block, pos token.Pos, lhs, rhs []ast.Expr,
 			stmts = append(stmts, ss...)
 
 			if r[0].Type == shaderir.NumberExpr {
-				switch block.vars[l[0].Index].typ.Main {
+				t, ok := block.findLocalVariableByIndex(l[0].Index)
+				if !ok {
+					cs.addError(pos, fmt.Sprintf("unexpected local variable index: %d", l[0].Index))
+					return nil, false
+				}
+				switch t.Main {
 				case shaderir.Int:
 					r[0].ConstType = shaderir.ConstTypeInt
 				case shaderir.Float:
diff --git a/internal/shader/testdata/blocks2.expected.vs b/internal/shader/testdata/blocks2.expected.vs
new file mode 100644
index 000000000..cc8c008bb
--- /dev/null
+++ b/internal/shader/testdata/blocks2.expected.vs
@@ -0,0 +1,12 @@
+void F0(out vec2 l0) {
+	bool l1 = false;
+	l1 = true;
+	{
+		int l2 = 0;
+		l2 = 0;
+		l0 = vec2(l2);
+		return;
+	}
+	l0 = vec2(1.0);
+	return;
+}
diff --git a/internal/shader/testdata/blocks2.go b/internal/shader/testdata/blocks2.go
new file mode 100644
index 000000000..d9f4c3557
--- /dev/null
+++ b/internal/shader/testdata/blocks2.go
@@ -0,0 +1,10 @@
+package main
+
+func Foo() vec2 {
+	x := true
+	{
+		x := 0
+		return vec2(x)
+	}
+	return vec2(1)
+}