Implement capturing outer this in closures

Author: Frotty

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

@@ -401,20 +401,30 @@ 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) {
+            receiver = globalState.ensureObject((ImClassType) va.getReceiver().attrTyp(),
+                ((ILconstInt) receiverVal).getVal(), 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

@@ -351,6 +351,16 @@ public ILconstObject toObject(ILconst val) {
         throw new InterpreterException(this, "Value " + val + " (" + val.getClass().getSimpleName() + ") cannot be cast to object.");
     }
 
+    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 static class StackTrace {
         private final List<ILStackFrame> stackFrames;
 

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

@@ -244,23 +244,93 @@ public ImType case_ImTypeVarRef(ImTypeVarRef t) {
      * class variables instead
      */
     private void transformTranslated(ImExpr t) {
-        final List<ImVarAccess> vas = Lists.newArrayList();
+        final List<ImVarAccess> capturedLocals = Lists.newArrayList();
+        final List<ImMemberAccess> capturedFields = Lists.newArrayList();
+
         t.accept(new ImExpr.DefaultVisitor() {
             @Override
             public void visit(ImVarAccess va) {
                 super.visit(va);
                 if (isLocalToOtherFunc(va.getVar())) {
-                    vas.add(va);
+                    capturedLocals.add(va);
+                }
+            }
+
+            @Override
+            public void visit(ImMemberAccess ma) {
+                super.visit(ma);
+                if (isCapturedOuterField(ma)) {
+                    capturedFields.add(ma);
                 }
             }
 
+            private boolean isCapturedOuterField(ImMemberAccess ma) {
+                // receiver must be *this* of the closure
+                if (!(ma.getReceiver() instanceof ImVarAccess)) {
+                    return false;
+                }
+                ImVarAccess recv = (ImVarAccess) ma.getReceiver();
+                if (recv.getVar() != tr.getThisVar(e)) {
+                    return false;
+                }
 
+                // field must belong to some class
+                var pp = ma.getVar().getParent();
+                if (pp == null || !(pp.getParent() instanceof ImClass)) {
+                    return false;
+                }
+
+                ImClass owner = (ImClass) pp.getParent();
+                // and it must not be a field of the closure class itself
+                return owner != c;
+            }
         });
 
-        for (ImVarAccess va : vas) {
+        // Existing behaviour: capture locals from outer functions
+        for (ImVarAccess va : capturedLocals) {
             ImVar v = getClosureVarFor(va.getVar());
-            va.replaceBy(JassIm.ImMemberAccess(e, closureThis(), JassIm.ImTypeArguments(), v, JassIm.ImExprs()));
+            va.replaceBy(JassIm.ImMemberAccess(e, closureThis(),
+                JassIm.ImTypeArguments(), v, JassIm.ImExprs()));
+        }
+
+        // New behaviour: capture outer "this" for fields of the enclosing class
+        for (ImMemberAccess ma : capturedFields) {
+            ImVar field = ma.getVar();
+            ImClass owner = (ImClass) field.getParent().getParent();
+
+            ImVar outerThis = findOuterThisVar(owner);
+            if (outerThis == null) {
+                // no obvious owning instance – leave it, fail later with a clear error if needed
+                continue;
+            }
+
+            // Create/get the closure field to store outer "this"
+            ImVar capturedOuterThisField = getClosureVarFor(outerThis);
+
+            // Rewrite receiver: this.capturedOuterThisField
+            ImExpr newReceiver = JassIm.ImMemberAccess(
+                e,
+                closureThis(),
+                JassIm.ImTypeArguments(),
+                capturedOuterThisField,
+                JassIm.ImExprs()
+            );
+            ma.setReceiver(newReceiver);
+        }
+    }
+
+    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,53 @@ 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()"
+        );
+    }
+
+
 
 }