Implement capturing outer this in closures (#1115)

* Implement capturing outer this in closures

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/intermediatelang/interpreter/EvaluateExpr.java

@@ -401,20 +401,33 @@ public ILconst get() {
 
     public static ILaddress evaluateLvalue(ImMemberAccess va, ProgramState globalState, LocalState localState) {
         ImVar v = va.getVar();
-        ILconstObject receiver = globalState.toObject(va.getReceiver().evaluate(globalState, localState));
+        ILconst receiverVal = va.getReceiver().evaluate(globalState, localState);
+        ILconstObject receiver = globalState.toObject(receiverVal);
+        if (receiver == null && receiverVal instanceof ILconstInt && va.getReceiver().attrTyp() instanceof ImClassType) {
+            int objectId = ((ILconstInt) receiverVal).getVal();
+            if (objectId != 0) {
+                receiver = globalState.ensureObject((ImClassType) va.getReceiver().attrTyp(),
+                    objectId, va.attrTrace());
+            }
+        }
+        if (receiver == null) {
+            throw new InterpreterException(va.attrTrace(), "Null pointer dereference");
+        }
+        ILconstObject receiverFinal = receiver;
         List<Integer> indexes =
             va.getIndexes().stream()
                 .map(ie -> ((ILconstInt) ie.evaluate(globalState, localState)).getVal())
                 .collect(Collectors.toList());
+        List<Integer> indexesFinal = indexes;
         return new ILaddress() {
             @Override
             public void set(ILconst value) {
-                receiver.set(v, indexes, value);
+                receiverFinal.set(v, indexesFinal, value);
             }
 
             @Override
             public ILconst get() {
-                return receiver.get(v, indexes)
+                return receiverFinal.get(v, indexesFinal)
                     .orElseGet(() -> va.attrTyp().defaultValue());
             }
         };

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/intermediatelang/interpreter/ProgramState.java

@@ -342,6 +342,16 @@ public ILconst getHandleByIndex(int val) {
         return handleMap.get(val);
     }
 
+    public ILconstObject ensureObject(ImClassType clazz, int objectId, Element trace) {
+        ILconstObject existing = indexToObject.get(objectId);
+        if (existing != null) {
+            return existing;
+        }
+        ILconstObject res = new ILconstObject(clazz, objectId, trace);
+        indexToObject.put(objectId, res);
+        return res;
+    }
+
     public ILconstObject toObject(ILconst val) {
         if (val instanceof ILconstObject) {
             return (ILconstObject) val;

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtranslation/ClosureTranslator.java

@@ -30,6 +30,7 @@ public class ClosureTranslator {
     private Map<ImTypeVar, ImTypeVar> typeVars;
     private ImFunction impl;
     private ImClass c;
+    private ImVar capturedThisField;
 
     public ClosureTranslator(ExprClosure e, ImTranslator tr, ImFunction f) {
         super();
@@ -167,6 +168,7 @@ private ImClass createClass() {
         } else {
             impl.getBody().add(JassIm.ImReturn(e, translated));
         }
+        captureEnclosingThis();
         transformTranslated(translated);
 
         typeVars = rewriteTypeVars(c);
@@ -245,6 +247,9 @@ public ImType case_ImTypeVarRef(ImTypeVarRef t) {
      */
     private void transformTranslated(ImExpr t) {
         final List<ImVarAccess> vas = Lists.newArrayList();
+        final List<ImMemberAccess> receiversToRewrite = Lists.newArrayList();
+        ImVar closureThisVar = tr.getThisVar(e);
+
         t.accept(new ImExpr.DefaultVisitor() {
             @Override
             public void visit(ImVarAccess va) {
@@ -254,13 +259,77 @@ public void visit(ImVarAccess va) {
                 }
             }
 
+            @Override
+            public void visit(ImMemberAccess ma) {
+                super.visit(ma);
+                if (capturedThisField != null && ma.getReceiver() instanceof ImVarAccess) {
+                    ImVar recvVar = ((ImVarAccess) ma.getReceiver()).getVar();
+                    if (recvVar == closureThisVar && capturesFromEnclosingThis(owningClass(ma.getVar()))) {
+                        receiversToRewrite.add(ma);
+                    }
+                }
+            }
 
         });
 
         for (ImVarAccess va : vas) {
             ImVar v = getClosureVarFor(va.getVar());
             va.replaceBy(JassIm.ImMemberAccess(e, closureThis(), JassIm.ImTypeArguments(), v, JassIm.ImExprs()));
         }
+
+        for (ImMemberAccess ma : receiversToRewrite) {
+            ma.setReceiver(JassIm.ImMemberAccess(e, closureThis(), JassIm.ImTypeArguments(), capturedThisField, JassIm.ImExprs()));
+        }
+    }
+
+    private void captureEnclosingThis() {
+        ImVar outerThis = getEnclosingThisVar();
+        if (outerThis != null) {
+            capturedThisField = getClosureVarFor(outerThis);
+        }
+    }
+
+    private ImVar getEnclosingThisVar() {
+        if (f != null && !f.getParameters().isEmpty()) {
+            ImVar param = f.getParameters().get(0);
+            if ("this".equals(param.getName()) && param.getType() instanceof ImClassType) {
+                return param;
+            }
+        }
+        return null;
+    }
+
+    private ImClass owningClass(ImVar var) {
+        if (var.getParent() != null && var.getParent().getParent() instanceof ImClass) {
+            return (ImClass) var.getParent().getParent();
+        }
+        return null;
+    }
+
+    private boolean capturesFromEnclosingThis(ImClass owner) {
+        if (owner == null || capturedThisField == null) {
+            return false;
+        }
+        if (!(capturedThisField.getType() instanceof ImClassType)) {
+            return false;
+        }
+
+        ImClass capturedClass = ((ImClassType) capturedThisField.getType()).getClassDef();
+        return capturedClass == owner || capturedClass.isSubclassOf(owner);
+    }
+
+    private ImVar findOuterThisVar(ImClass owner) {
+        // in instance methods, the first parameter is typically "this"
+        for (ImVar p : f.getParameters()) {
+            ImType t = p.getType();
+            if (t instanceof ImClassType) {
+                ImClassType ct = (ImClassType) t;
+                if (ct.getClassDef() == owner) {
+                    return p;
+                }
+            }
+        }
+        return null;
     }
 
     private ImVarAccess closureThis() {

de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/GenericsWithTypeclassesTests.java

@@ -1531,5 +1531,88 @@ public void mixingLegacyOwner_newType_insideGenericClassMethod() {
         );
     }
 
+    @Test
+    public void arrayListCapacity_lazyClosure_repro() {
+        testAssertOkLines(true,
+            "package test",
+            "",
+            "native testSuccess()",
+            "",
+            "constant int JASS_MAX_ARRAY_SIZE = 8190",
+            "",
+            "public function lazy<T:>(Lazy<T> l) returns Lazy<T>",
+            "    return l",
+            "",
+            "public abstract class Lazy<T:>",
+            "    T val = null",
+            "    boolean wasRetrieved = false",
+            "",
+            "    abstract function retrieve() returns T",
+            "",
+            "    function get() returns T",
+            "        if not wasRetrieved",
+            "            val = retrieve()",
+            "            wasRetrieved = true",
+            "        return val",
+            "",
+            "public class ArrayList<T:>",
+            "",
+            "public class CFBuilding",
+            "    ArrayList<CFBuilding> upgrades = null",
+            "    Lazy<boolean> hasAAUpgrade = lazy<boolean>(() -> begin",
+            "            var result = false",
+            "            if upgrades != null",
+        "                   result = true",
+            "            return result",
+            "        end)",
+            "",
+            "",
+            "init",
+            "    let a = new CFBuilding()",
+            "    let b = new CFBuilding()",
+            "    // ensure upgrades is non-null so closure does some work",
+            "    b.upgrades = new ArrayList<CFBuilding>()",
+            "    // invoke lazy attribute so its closure is actually referenced",
+            "    if b.hasAAUpgrade.get() and not a.hasAAUpgrade.get()",
+            "        testSuccess()"
+        );
+    }
+
+    @Test
+    public void inheritedField_lazyClosure_uses_enclosing_receiver() {
+        testAssertOkLines(true,
+            "package test",
+            "",
+            "native testSuccess()",
+            "",
+            "public abstract class Lazy<T:>",
+            "    T val = null",
+            "    boolean wasRetrieved = false",
+            "",
+            "    abstract function retrieve() returns T",
+            "",
+            "    function get() returns T",
+            "        if not wasRetrieved",
+            "            val = retrieve()",
+            "            wasRetrieved = true",
+            "        return val",
+            "",
+            "public function lazy<T:>(Lazy<T> l) returns Lazy<T>",
+            "    return l",
+            "",
+            "public class BaseBuilding",
+            "    boolean hasDetector = true",
+            "",
+            "public class AdvancedBuilding extends BaseBuilding",
+            "    Lazy<boolean> detectorAvailable = lazy<boolean>(() -> hasDetector)",
+            "",
+            "init",
+            "    let b = new AdvancedBuilding()",
+            "    b.hasDetector = true",
+            "    if b.detectorAvailable.get()",
+            "        testSuccess()"
+        );
+    }
+
 
 }