Track locals, globals, and upvalues in ast

src/jse/org/luaj/vm2/ast/Block.java

@@ -27,6 +27,7 @@ import java.util.List;
 public class Block extends Stat {
 	
 	public List<Stat> stats = new ArrayList<Stat>();
+	public NameScope scope;
 	
 	public void add(Stat s) {
 		if ( s == null )

src/jse/org/luaj/vm2/ast/Exp.java

@@ -106,6 +106,8 @@ public class Exp {
 		public boolean isvarexp() {
 			return true;
 		}
+		public void markHasAssignment() {
+		}
 	}
 	
 	public static class NameExp extends VarExp {
@@ -113,7 +115,9 @@ public class Exp {
 		public NameExp(String name) {
 			this.name = new Name(name);
 		}
-		
+		public void markHasAssignment() {
+			name.variable.hasassignments = true;
+		}
 		public void accept(Visitor visitor) {
 			visitor.visit(this);
 		}

src/jse/org/luaj/vm2/ast/FuncBody.java

@@ -24,6 +24,7 @@ package org.luaj.vm2.ast;
 public class FuncBody {
 	public ParList parlist;
 	public Block block;
+	public NameScope scope;
 
 	public FuncBody(ParList parlist, Block block) {
 		this.parlist = parlist!=null? parlist: ParList.EMPTY_PARLIST;

src/jse/org/luaj/vm2/ast/FuncName.java

@@ -28,7 +28,7 @@ public class FuncName {
 	// example: a.b.c.d:e
 	
 	// initial base name: "a"
-	public final String name;
+	public final Name name;
 	
 	// intermediate field accesses: "b", "c", "d"
 	public List<String> dots;
@@ -37,7 +37,7 @@ public class FuncName {
 	public String method;
 	
 	public FuncName( String name ) {
-		this.name = name;
+		this.name = new Name(name);
 	}
 	
 	public void adddot(String dot) {

src/jse/org/luaj/vm2/ast/Name.java

@@ -21,8 +21,11 @@
 ******************************************************************************/
 package org.luaj.vm2.ast;
 
+import org.luaj.vm2.ast.NameScope.NamedVariable;
+
 public class Name {
 	public final String name;
+	public NamedVariable variable;
 	public Name(String name) {
 		this.name = name;
 	}

src/jse/org/luaj/vm2/ast/NameResolver.java

@@ -0,0 +1,98 @@
+package org.luaj.vm2.ast;
+
+import java.util.List;
+
+import org.luaj.vm2.ast.Exp.NameExp;
+import org.luaj.vm2.ast.Exp.VarExp;
+import org.luaj.vm2.ast.NameScope.NamedVariable;
+import org.luaj.vm2.ast.Stat.Assign;
+import org.luaj.vm2.ast.Stat.GenericFor;
+import org.luaj.vm2.ast.Stat.LocalAssign;
+import org.luaj.vm2.ast.Stat.LocalFuncDef;
+import org.luaj.vm2.ast.Stat.NumericFor;
+
+/** 
+ * Visitor that resolves names to scopes.
+ * Each Name is resolved to a NamedVarible, possibly in a NameScope 
+ * if it is a local, or in no named scope if it is a global. 
+ */
+public class NameResolver extends Visitor {
+
+	private NameScope scope = null;
+
+	private void pushScope() {
+		scope = new NameScope(scope);
+	}
+	private void popScope() {
+		scope = scope.outerScope;
+	}
+	
+	public void visit(NameScope scope) {
+	}	
+
+	public void visit(Block block) {
+		pushScope();
+		block.scope = scope;
+		super.visit(block);
+		popScope();
+	}
+	
+	public void visit(FuncBody body) {
+		pushScope();
+		scope.functionNestingCount++;
+		body.scope = scope;
+		for ( Name n : body.parlist.names )
+			defineLocalVar(n);
+		super.visit(body);
+		popScope();
+	}
+
+	public void visit(LocalFuncDef stat) {
+		defineLocalVar(stat.name);
+		super.visit(stat);
+	}
+
+	public void visit(NumericFor stat) {
+		pushScope();
+		stat.scope = scope;
+		defineLocalVar(stat.name);
+		super.visit(stat);
+		popScope();
+	}
+
+	public void visit(GenericFor stat) {
+		pushScope();
+		stat.scope = scope;
+		defineLocalVars( stat.names );
+		super.visit(stat);
+		popScope();
+	}
+
+	public void visit(NameExp exp) {
+		NamedVariable v = scope.find(exp.name.name);
+		exp.name.variable = v;
+		if ( v.isLocal() && scope.functionNestingCount != v.definingScope.functionNestingCount )
+			v.isupvalue = true;
+		super.visit(exp);
+	}
+	
+	public void visit(Assign stat) {
+		super.visit(stat);
+		for ( VarExp v : stat.vars )
+			v.markHasAssignment();
+	}
+
+	public void visit(LocalAssign stat) {
+		defineLocalVars( stat.names );
+		super.visit(stat);
+	}
+
+	protected void defineLocalVars(List<Name> names) {
+		for ( Name n : names ) 
+			defineLocalVar(n);
+	}
+
+	protected void defineLocalVar(Name name) {
+		name.variable = scope.define(name.name);
+	}
+}

src/jse/org/luaj/vm2/ast/NameScope.java

@@ -0,0 +1,84 @@
+package org.luaj.vm2.ast;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class NameScope {
+
+	private static final Set<String> LUA_KEYWORDS = new HashSet<String>();
+	
+	static { 
+		String[] k = new String[] { 
+	        "and", "break", "do", "else", "elseif", "end", 
+	        "false", "for", "function", "if", "in", "local",
+	        "nil", "not", "or", "repeat", "return", 
+	        "then", "true", "until", "while" };
+		for ( int i=0; i<k.length; i++ )
+			LUA_KEYWORDS.add( k[i] );
+	}
+	
+	public final Map<String,NamedVariable> namedVariables = new HashMap<String,NamedVariable>();
+
+	public final NameScope outerScope;
+
+	public int functionNestingCount;
+	
+	/** Construct default names scope */
+	public NameScope() {
+		this.outerScope = null;
+		this.functionNestingCount = 0;
+	}
+	
+	/** Construct name scope within another scope*/
+	public NameScope(NameScope outerScope) {
+		this.outerScope = outerScope;
+		this.functionNestingCount = outerScope!=null? outerScope.functionNestingCount: 0;
+	}
+	
+	/** Look up a name.  If it is a global name, then throw IllegalArgumentException. */
+	public NamedVariable find( String name ) throws IllegalArgumentException {
+		validateIsNotKeyword(name);
+		for ( NameScope n = this; n!=null; n=n.outerScope )
+			if ( n.namedVariables.containsKey(name) )
+				return namedVariables.get(name);
+		NamedVariable value = new NamedVariable(name);
+		this.namedVariables.put(name, value);
+		return value;
+	}
+	
+	/** Define a name in this scope.  If it is a global name, then throw IllegalArgumentException. */
+	public NamedVariable define( String name ) throws IllegalStateException, IllegalArgumentException {
+		validateIsNotKeyword(name);
+		NamedVariable value = new NamedVariable(name, this);
+		this.namedVariables.put(name, value);
+		return value;
+	}
+	
+	private void validateIsNotKeyword(String name) {
+		if ( LUA_KEYWORDS.contains(name) )
+			throw new IllegalArgumentException("name is a keyword: '"+name+"'");
+	}
+
+	/** Named variable is load, global, or upvalue.  */
+	public static class NamedVariable {
+		public final String name;
+		public final NameScope definingScope;
+		public boolean isupvalue;
+		public boolean hasassignments;
+		/** Global is named variable not associated with a defining scope */
+		public NamedVariable(String name) {
+			this.name = name;
+			this.definingScope = null;
+		}
+		public NamedVariable(String name, NameScope definingScope) {
+		/** Local variable is defined in a particular scope.  */
+			this.name = name;
+			this.definingScope = definingScope;
+		}
+		public boolean isLocal() {
+			return this.definingScope != null;
+		}
+	}
+}

src/jse/org/luaj/vm2/ast/Stat.java

@@ -23,6 +23,8 @@ package org.luaj.vm2.ast;
 
 import java.util.List;
 
+import org.luaj.vm2.ast.Exp.NameExp;
+
 abstract
 public class Stat {
 	public abstract void accept(Visitor visitor);
@@ -176,6 +178,7 @@ public class Stat {
 		public List<Name> names;
 		public List<Exp> exps;
 		public Block block;
+		public NameScope scope;
 		public GenericFor(List<Name> names, List<Exp> exps, Block block) {
 			this.names = names;
 			this.exps = exps;
@@ -191,6 +194,7 @@ public class Stat {
 		public final Name name;
 		public final Exp initial,limit,step;
 		public final Block block;
+		public NameScope scope;
 		public NumericFor(String name, Exp initial, Exp limit, Exp step, Block block) {
 			this.name = new Name(name);
 			this.initial = initial;

src/jse/org/luaj/vm2/ast/Visitor.java

@@ -7,6 +7,7 @@ abstract public class Visitor {
 		chunk.block.accept(this); 
 	};
 	public void visit(Block block) {
+		visit(block.scope);
 		if ( block.stats != null )
 			for ( Stat s: block.stats )
 				s.accept(this);
@@ -24,6 +25,7 @@ abstract public class Visitor {
 		stat.body.accept(this);
 	}
 	public void visit(Stat.GenericFor stat) {
+		visit(stat.scope);
 		visitNames(stat.names);
 		visitExps(stat.exps);
 		stat.block.accept(this);
@@ -46,6 +48,7 @@ abstract public class Visitor {
 		stat.body.accept(this);
 	}
 	public void visit(Stat.NumericFor stat) {
+		visit(stat.scope);
 		visit(stat.name);
 		stat.initial.accept(this);
 		stat.limit.accept(this);
@@ -65,6 +68,7 @@ abstract public class Visitor {
 		stat.block.accept(this);
 	}
 	public void visit(FuncBody body) {
+		visit(body.scope);
 		body.parlist.accept(this);
 		body.block.accept(this);
 	}
@@ -138,8 +142,10 @@ abstract public class Visitor {
 			for ( Name n: names )
 				visit(n);
 	}
-	private void visit(Name name) {
+	public void visit(Name name) {
 	}
-	private void visit(String name) {
+	public void visit(String name) {
+	}
+	public void visit(NameScope scope) {
 	}
 }