(fix) Throw exception when assignment to captured variables, issue #101

src/main/java/com/googlecode/aviator/runtime/function/LambdaFunction.java

@@ -1,7 +1,10 @@
 package com.googlecode.aviator.runtime.function;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import com.googlecode.aviator.BaseExpression;
 import com.googlecode.aviator.Expression;
 import com.googlecode.aviator.runtime.type.AviatorObject;
 import com.googlecode.aviator.utils.Env;
@@ -26,7 +29,14 @@ public abstract class LambdaFunction extends AbstractFunction {
   public LambdaFunction(List<String> arguments, Expression expression, Env context) {
     super();
     this.arguments = arguments;
-    this.context = context.clone();
+    this.context = context;
+    Set<String> argumentSet = new HashSet<>(this.arguments);
+    for (String var : expression.getVariableNames()) {
+      if (!var.contains(".") && !argumentSet.contains(var)) {
+        // mark the var is captured.
+        context.capture(var, ((BaseExpression) expression).getExpression());
+      }
+    }
     this.expression = expression;
   }
 

src/main/java/com/googlecode/aviator/runtime/op/OperationRuntime.java

@@ -166,7 +166,7 @@ private static void trace(Map<String, Object> env, OperatorType opType, AviatorO
       case 3:
         RuntimeUtils.printTrace(env,
             TRACE_PREFIX + args[0].desc(env) + WHITE_SPACE + "?" + WHITE_SPACE + args[0].desc(env)
-            + WHITE_SPACE + ":" + WHITE_SPACE + args[1].desc(env) + " => " + result.desc(env));
+                + WHITE_SPACE + ":" + WHITE_SPACE + args[1].desc(env) + " => " + result.desc(env));
         break;
     }
   }

src/main/java/com/googlecode/aviator/runtime/type/AviatorJavaType.java

@@ -259,6 +259,7 @@ public AviatorObject setValue(AviatorObject value, Map<String, Object> env) {
     if (this.name.contains(".")) {
       throw new IllegalArgumentException("Can't assignment value to `" + this.name + "`");
     }
+
     Object v = value.getValue(env);
     env.put(this.name, v);
     return new AviatorRuntimeJavaType(v);

src/main/java/com/googlecode/aviator/utils/Env.java

@@ -31,6 +31,7 @@
 import java.util.Map;
 import java.util.Set;
 import com.googlecode.aviator.AviatorEvaluatorInstance;
+import com.googlecode.aviator.exception.ExpressionRuntimeException;
 import com.googlecode.aviator.lexer.token.Variable;
 
 /**
@@ -68,21 +69,22 @@ public void setInstance(AviatorEvaluatorInstance instance) {
 
   public static final Map<String, Object> EMPTY_ENV = Collections.emptyMap();
 
+  private Map<String, String> capturedVars;
+
+  public void capture(String var, String expression) {
+    if (capturedVars == null) {
+      capturedVars = new HashMap<>();
+    }
+    capturedVars.put(var, expression);
+  }
+
   /**
    * Constructs an env instance with empty state.
    */
   public Env() {
     this(EMPTY_ENV);
   }
 
-  @Override
-  public Env clone() {
-    Env ret = new Env(this.mDefaults == EMPTY_ENV ? EMPTY_ENV : new HashMap<>(this.mDefaults),
-        new HashMap<>(this.mOverrides));
-    ret.instance = this.instance;
-    return ret;
-  }
-
   /**
    * Constructor.
    *
@@ -210,6 +212,12 @@ public Set<String> keySet() {
    */
   @Override
   public Object put(String key, Object value) {
+    String capturedExp = null;
+    if (this.capturedVars != null && this.capturedVars.containsKey(key)) {
+      throw new ExpressionRuntimeException("Can't assignment value to captured variable.The `" + key
+          + "` is already captured by lambda.");
+    }
+
     Object prior;
     Map<String, Object> overrides = getmOverrides(false);
     if (overrides.containsKey(key)) {
@@ -282,8 +290,8 @@ public Collection<Object> values() {
   public String toString() {
     StringBuffer buf = new StringBuffer(32 * size());
     buf.append(super.toString()).append("{"). //
-    append(Variable.INSTANCE_VAR).append("=").append(this.instance).append(", ").//
-    append(Variable.ENV_VAR).append("=").append("<this>");
+        append(Variable.INSTANCE_VAR).append("=").append(this.instance).append(", ").//
+        append(Variable.ENV_VAR).append("=").append("<this>");
 
     Iterator<String> it = keySet().iterator();
     boolean hasNext = it.hasNext();

src/test/java/com/googlecode/aviator/LambdaUnitTest.java

@@ -57,6 +57,18 @@ public void testLambdaClosure() {
     assertEquals(19, AviatorEvaluator.execute("test(4)(5)(6) + a", env));
   }
 
+  @Test(expected = ExpressionRuntimeException.class)
+  public void testIssue101() {
+    String exp = "a=1; b = lambda(x) -> a+ x end ; a=4 ; b(5)";
+    AviatorEvaluator.execute(exp); // output 6
+  }
+
+  @Test
+  public void testIssue101_assignment() {
+    String exp = "a=1; c=a; b = lambda(x) -> a+ x end ; c=4 ; b(5)";
+    assertEquals(6, AviatorEvaluator.execute(exp));
+  }
+
   public static class Foo {
     private int a;
 

src/test/java/com/googlecode/aviator/example/LambdaExample.java

@@ -1,14 +1,27 @@
 package com.googlecode.aviator.example;
 
+import java.util.HashMap;
+import java.util.Map;
 import com.googlecode.aviator.AviatorEvaluator;
 
 public class LambdaExample {
 
   public static void main(String[] args) {
 
-    String exp = "a=1; b = lambda(x) -> a+ x end ; a=4 ; b(5)";
+    String exp = "a=1; b = lambda(x) -> a+ x end ; b(5)";
     System.out.println(AviatorEvaluator.execute(exp)); // output 6
 
 
+    Map<String, Object> env = new HashMap<String, Object>();
+    env.put("x", 1);
+    env.put("y", 2);
+    env.put("z", 3);
+
+    AviatorEvaluator.defineFunction("test",
+        "lambda(x) -> lambda(y) -> lambda(z) -> x + y + z end end end");
+    System.out.println(AviatorEvaluator.execute("test(4)(5)(6)", env)); // output 15
+
+    env.put("a", 4);
+    System.out.println(AviatorEvaluator.execute("test(4)(5)(6) + a", env)); // output 19
   }
 }