From 576cf2068dcf3f6482089fe16f9530074feb320b Mon Sep 17 00:00:00 2001
From: James Roseborough <james.roseborough@luaj.org>
Date: Wed, 14 Apr 2010 05:36:43 +0000
Subject: [PATCH] Special processing on for loop upvalues.

---
 src/jse/org/luaj/vm2/luajc/JavaBuilder.java | 17 +++++++++
 src/jse/org/luaj/vm2/luajc/JavaGen.java     | 22 ++++++++----
 src/jse/org/luaj/vm2/luajc/Slots.java       | 38 +++++++++++++++++++++
 3 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/src/jse/org/luaj/vm2/luajc/JavaBuilder.java b/src/jse/org/luaj/vm2/luajc/JavaBuilder.java
index 600edcf8..3a96a854 100644
--- a/src/jse/org/luaj/vm2/luajc/JavaBuilder.java
+++ b/src/jse/org/luaj/vm2/luajc/JavaBuilder.java
@@ -338,6 +338,23 @@ public class JavaBuilder {
 			append(new ASTORE(index));
 		}
 	}
+
+	public void createUpvalues(int pc, int firstslot, int numslots) {
+		for ( int i=0; i<numslots; i++ ) {
+			int slot = firstslot + i;
+			boolean isupcreate = slots.isUpvalueCreate(pc, slot);
+			if ( isupcreate ) {
+				int index = findSlotIndex( slot, true );
+				append(new PUSH(cp, 1));
+				append(new ANEWARRAY(cp.addClass(STR_LUAVALUE)));
+				dup();
+				append(new PUSH(cp, 0));
+				loadNil();
+				append(InstructionConstants.AASTORE);
+				append(new ASTORE(index));
+			}
+		}
+	}
 	
 	private static String upvalueName(int upindex) {
 		return PREFIX_UPVALUE+upindex;
diff --git a/src/jse/org/luaj/vm2/luajc/JavaGen.java b/src/jse/org/luaj/vm2/luajc/JavaGen.java
index 0562683d..5de0d8d6 100644
--- a/src/jse/org/luaj/vm2/luajc/JavaGen.java
+++ b/src/jse/org/luaj/vm2/luajc/JavaGen.java
@@ -175,8 +175,14 @@ public class JavaGen {
 				break;
 				
 			case Lua.OP_JMP: /*	sBx	pc+=sBx					*/
+			{
+				int pc1 = pc+1+sbx;
+				ins = p.code[pc1];
+				if ( Lua.GET_OPCODE(ins) == Lua.OP_TFORLOOP )
+					builder.createUpvalues(pc, Lua.GETARG_A(ins)+3, Lua.GETARG_C(ins));
 				builder.addBranch(pc, JavaBuilder.BRANCH_GOTO, pc+1+sbx);
 				break;
+			}
 				
 			case Lua.OP_EQ: /*	A B C	if ((RK(B) == RK(C)) ~= A) then pc++		*/
 			case Lua.OP_LT: /*	A B C	if ((RK(B) <  RK(C)) ~= A) then pc++  		*/
@@ -280,13 +286,15 @@ public class JavaGen {
 				break;
 				
 			case Lua.OP_RETURN: /*	A B	return R(A), ... ,R(A+B-2)	(see note)	*/
-				if ( c == 1 )
-					break;
-				switch ( b ) {
-				case 0: loadVarargResults( builder, pc, a, vresultbase ); break;
-				case 1: builder.loadNone(); break;
-				case 2: builder.loadLocal(pc, a); break;
-				default: builder.newVarargs(pc, a, b-1); break;
+				if ( c == 1 ) {
+					builder.loadNone();
+				} else {
+					switch ( b ) {
+					case 0: loadVarargResults( builder, pc, a, vresultbase ); break;
+					case 1: builder.loadNone(); break;
+					case 2: builder.loadLocal(pc, a); break;
+					default: builder.newVarargs(pc, a, b-1); break;
+					}
 				}
 				builder.areturn(); 
 				break;
diff --git a/src/jse/org/luaj/vm2/luajc/Slots.java b/src/jse/org/luaj/vm2/luajc/Slots.java
index f1941634..8eedfedc 100644
--- a/src/jse/org/luaj/vm2/luajc/Slots.java
+++ b/src/jse/org/luaj/vm2/luajc/Slots.java
@@ -90,6 +90,7 @@ public class Slots {
 		markassignments( p );
 		markuninitialized( p );
 		markupvalues( p );
+		markforloopupvalues( p );
 	}
 	
 	public String toString() {
@@ -329,6 +330,43 @@ public class Slots {
 		}
 	}
 	
+	private void markforloopupvalues( Prototype p ) {
+		for ( int pc1=n; --pc1>=0; ) {
+			int i = p.code[pc1];
+			if ( Lua.GET_OPCODE(i) == Lua.OP_TFORLOOP ) {
+				int a = Lua.GETARG_A(i);
+				int c = Lua.GETARG_C(i);
+				for ( int pc0=pc1; --pc0>=0; ) {
+					i = p.code[pc0];
+					int o = Lua.GET_OPCODE(i);
+					int sbx = Lua.GETARG_sBx(i);
+					if ( o == Lua.OP_JMP && (pc0 + 1 + sbx == pc1) ) {
+						for ( int j=0; j<c; j++ ) {
+							checkPromoteLoopUpvalue( pc0, pc1, a+3+j );
+						}
+					}
+				}
+			}
+		}
+		
+	}
+	
+	private void checkPromoteLoopUpvalue(int pc0, int pc1, int slot) {
+		for ( int index=pc0+1; index<=pc1; ++index ) {
+			switch (slots[index][slot]) {
+			case UPVAL_CREATE:
+			case UPVAL_USE_CREATE:
+			case UPVAL_USE_ASSIGN:
+			case UPVAL_USE:
+				int i = pc0;
+				slots[++i][slot] = UPVAL_CREATE;
+				while ( ++i<=pc1 )
+					slots[i][slot] = UPVAL_USE;
+				return;
+			}
+		}
+	}
+
 	private void promoteUpvalueBefore(byte[] s, int index, int j) {
 		int begin  = prevUndefined(index,j);
 		int assign = firstAssignAfter(begin,index,j);