Add check to handle recursive definitions (like "x=[x]") correctly

cs/AST.cs

@@ -75,6 +75,14 @@ public abstract class ASTNode {
 	// Returns a simplified version of this node (may be a new node, or this node unchanged)
 	public abstract ASTNode Simplify();
 
+	// Does this node contain the variable `variableName`?
+	// This is used to determine whether the Value of an AssignmentNode is defined in terms
+	// of its Variable, in which case a temporary register needs to be used during
+	// calculations.
+	public virtual Boolean ContainsVariable(String variableName) {
+		return false;
+	}
+
 	// Copy the source line from this node to the given node and return it.
 	// Call as `return CopyLine(new SomeNode(...));` inside Simplify() overrides.
 	protected ASTNode CopyLine(ASTNode result) {
@@ -144,6 +152,10 @@ public override ASTNode Simplify() {
 		return this;
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Name == variableName;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -232,6 +244,10 @@ public override ASTNode Simplify() {
 		return new UnaryOpNode(Op, simplifiedOperand);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Operand.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -332,6 +348,11 @@ public override ASTNode Simplify() {
 		return new BinaryOpNode(Op, simplifiedLeft, simplifiedRight);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Left.ContainsVariable(variableName)
+			|| Right.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -364,6 +385,13 @@ public override ASTNode Simplify() {
 		return new ComparisonChainNode(simplifiedOperands, Operators);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		for (Int32 i = 0; i < Operands.Count; ++i) {
+			if (Operands[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -402,6 +430,14 @@ public override ASTNode Simplify() {
 		return CopyLine(new CallNode(Function, simplifiedArgs));
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		if (Function == variableName) return true;
+		for (Int32 i = 0; i < Arguments.Count; ++i) {
+			if (Arguments[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -425,6 +461,10 @@ public override ASTNode Simplify() {
 		return Expression.Simplify();
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Expression.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -460,6 +500,13 @@ public override ASTNode Simplify() {
 		return new ListNode(simplifiedElements);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		for (Int32 i = 0; i < Elements.Count; ++i) {
+			if (Elements[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -500,6 +547,14 @@ public override ASTNode Simplify() {
 		return new MapNode(simplifiedKeys, simplifiedValues);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		for (Int32 i = 0; i < Keys.Count; ++i) {
+			if (Keys[i].ContainsVariable(variableName)) return true;
+			if (Values[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -523,6 +578,11 @@ public override ASTNode Simplify() {
 		return new IndexNode(Target.Simplify(), Index.Simplify());
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Target.ContainsVariable(variableName)
+			|| Index.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -552,6 +612,12 @@ public override ASTNode Simplify() {
 		return new SliceNode(Target.Simplify(), simplifiedStart, simplifiedEnd);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Target.ContainsVariable(variableName)
+			|| StartIndex.ContainsVariable(variableName)
+			|| EndIndex.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -575,6 +641,10 @@ public override ASTNode Simplify() {
 		return new MemberNode(Target.Simplify(), Member);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		return Target.ContainsVariable(variableName);
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -611,6 +681,14 @@ public override ASTNode Simplify() {
 		return new MethodCallNode(Target.Simplify(), Method, simplifiedArgs);
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		if (Target.ContainsVariable(variableName)) return true;
+		for (Int32 i = 0; i < Arguments.Count; ++i) {
+			if (Arguments[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}
@@ -646,6 +724,14 @@ public override ASTNode Simplify() {
 		return CopyLine(new ExprCallNode(Function.Simplify(), simplifiedArgs));
 	}
 
+	public override Boolean ContainsVariable(String variableName) {
+		if (Function.ContainsVariable(variableName)) return true;
+		for (Int32 i = 0; i < Arguments.Count; ++i) {
+			if (Arguments[i].ContainsVariable(variableName)) return true;
+		}
+		return false;
+	}
+
 	public override Int32 Accept(IASTVisitor visitor) {
 		return visitor.Visit(this);
 	}

cs/CodeGenerator.cs

@@ -361,13 +361,36 @@ public Int32 Visit(AssignmentNode node) {
 			// the target register when done.  But we should probably return to
 			// this later and see if we can optimize it more.
 			varReg = AllocReg();
-			_variableRegs[node.Variable] = varReg;
+			// Don't set it in _variableRegs just yet.
+		}
+		// In the case of lists and maps, it's possible to define a variable which
+		// relies on the variable's prior value; for example, ensureImport from the
+		// importUtil library has the code:
+		// if moduleName isa string then moduleName = [moduleName]
+		// In these sorts of cases specifically, we need to allocate a temporary register
+		// for the new value of the variable and then copy the new value into the right
+		// register.
+		// Note that a similar case does not need to be added for IndexedAssignmentNode
+		// since that allocates a new register for the value in all cases.
+		// Holds the register for the variable in the recursive case (-1 = no need)
+		Int32 pendingReg = -1;
+		if (node.Value.ContainsVariable(node.Variable)) {
+			// Check and make sure this isn't something like "x = x"
+			// (we only need to check if the node is an IdentifierNode, since an
+			// IdentifierNode with any other name wouldn't trigger the recursive
+			// definition check)
+			IdentifierNode valueIdNode = node.Value as IdentifierNode;
+			if (valueIdNode == null) {
+				pendingReg = varReg;
+				varReg = AllocReg();
+			}
 		}
 		// Emit a NAME op unless one already dominates this point.  This must run
 		// on every path that assigns the variable, including conditional branches
 		// (e.g. the else clause of a single-line if), or the variable would be
 		// undefined at runtime when only that path executes.
-		EnsureNamed(node.Variable, varReg);
+		// Defer this line if we're doing a recursive definition.
+		if (pendingReg == -1) EnsureNamed(node.Variable, varReg);
 		// If the RHS is a function expression, note the current function count so we
 		// can assign the variable name to the resulting FuncDef afterward.
 		FunctionNode rhsFunc = node.Value as FunctionNode;
@@ -381,6 +404,22 @@ public Int32 Visit(AssignmentNode node) {
 			if (rhsFuncDef != null) rhsFuncDef.Name = node.Variable;
 		}
 
+		if (pendingReg != -1) {
+			// We allocated a temporary register earlier, so now we need to finish the
+			// assignment. Load the variable's actual register with the value of the
+			// temporary register and then free the temporary register to be used later.
+			EnsureNamed(node.Variable, pendingReg);
+			_emitter.EmitABC(Opcode.LOAD_rA_rB, pendingReg, varReg,
+				0, $"r{pendingReg} = r{varReg}");
+			FreeReg(varReg);
+			varReg = pendingReg;
+		}
+
+		// If the variable didn't exist before, now it does.
+		if (!_variableRegs.ContainsKey(node.Variable)) {
+			_variableRegs[node.Variable] = varReg;
+		}
+
 		// Note that we don't FreeReg(varReg) here, as we need this register to
 		// continue to serve as the storage for this variable for the life of
 		// the function.  (TODO: or until some lifetime analysis determines that

tests/testSuite.txt

@@ -1235,6 +1235,110 @@ n1 == n3: 0
 n1 == n4: 0
 n1 == n5: 0
 ================================================================================
+==== SECTION 8b: RECURSIVELY DEFINED VARIABLES
+================================================================================
+==== List containing itself
+================================
+x = 42
+x = [x]
+print x
+--------------------------------
+[42]
+================================
+==== Map containing itself
+================================
+x = "hello"
+x = {x: x}
+print x
+--------------------------------
+{"hello": "hello"}
+================================
+==== Operations A
+================================
+printAndReset = function
+	print outer.x
+	outer.x = 42
+end function
+x = 42
+printAndReset
+x = [-x]
+printAndReset
+x = [x * 2]
+printAndReset
+x = [(x)]
+printAndReset
+x = [[x]]
+printAndReset
+x = [{x: x}]
+printAndReset
+x = [0 <= x <= 50]
+printAndReset
+--------------------------------
+42
+[-42]
+[84]
+[42]
+[[42]]
+[{42: 42}]
+[1]
+================================
+==== Operations B
+================================
+printAndReset = function(v=42)
+	print outer.x
+	outer.x = @v
+end function
+exFunction = function(n)
+	return n*2
+end function
+exList = range(0,100)
+exMap = {42: "life, the universe and everything", "foo": "bar"}
+exMap.member = function(x)
+	return self[x]
+end function
+// Make sure printAndReset works
+x = 42
+// Function calls
+printAndReset @exFunction
+x = [x(2)]
+printAndReset
+x = [exFunction(x)]
+// Index/member
+printAndReset exMap
+x = [x[42]]
+printAndReset
+x = [exMap[x]]
+printAndReset exMap
+x = [x.foo]
+printAndReset exMap
+x = [x.member(42)]
+printAndReset
+x = [exMap.member(x)]
+// Expr call
+printAndReset exMap
+x = [x["member"](42)]
+// Slice
+printAndReset exList
+x = [x[5:10]]
+printAndReset
+x = [exList[x:50]]
+printAndReset
+x = [exList[40:x]]
+printAndReset
+--------------------------------
+42
+[4]
+[84]
+["life, the universe and everything"]
+["life, the universe and everything"]
+["bar"]
+["life, the universe and everything"]
+["life, the universe and everything"]
+["life, the universe and everything"]
+[[5, 6, 7, 8, 9]]
+[[42, 43, 44, 45, 46, 47, 48, 49]]
+[[40, 41]]
+================================================================================
 ==== SECTION 9: WHILE LOOPS
 ================================================================================
 ==== Basic while loop