HubSpot · andrewju · Apr 24, 2023 · Apr 24, 2023
diff --git a/src/main/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunction.java b/src/main/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunction.java
@@ -24,10 +24,14 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.StringJoiner;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import org.apache.commons.lang3.StringUtils;
 
 @Beta
 public class EagerMacroFunction extends MacroFunction {
@@ -212,6 +216,8 @@ public String reconstructImage() {
    */
   public String reconstructImage(String fullName) {
     String prefix = "";
+    StringBuilder result = new StringBuilder();
+    String setTagForAliasedVariables = getSetTagForAliasedVariables(fullName);
     String suffix = "";
     JinjavaInterpreter interpreter = JinjavaInterpreter.getCurrent();
     Optional<String> importFile = getImportFile(interpreter);
@@ -247,7 +253,6 @@ public String reconstructImage(String fullName) {
         );
     }
 
-    String result;
     if (
       (
         interpreter.getContext().getMacroStack().contains(getName()) &&
@@ -262,11 +267,19 @@ public String reconstructImage(String fullName) {
         String evaluation = (String) evaluate(
           getArguments().stream().map(arg -> DeferredMacroValueImpl.instance()).toArray()
         );
-        result =
-          (getStartTag(fullName, interpreter) + evaluation + getEndTag(interpreter));
+        result
+          .append(getStartTag(fullName, interpreter))
+          .append(setTagForAliasedVariables)
+          .append(evaluation)
+          .append(getEndTag(interpreter));
       } catch (DeferredValueException e) {
         // In case something not eager-supported encountered a deferred value
-        result = super.reconstructImage();
+        if (StringUtils.isNotEmpty(setTagForAliasedVariables)) {
+          throw new DeferredValueException(
+            "Aliased variables in not eagerly reconstructible macro function"
+          );
+        }
+        result.append(super.reconstructImage());
       } finally {
         reconstructing = false;
         interpreter
@@ -280,6 +293,26 @@ public String reconstructImage(String fullName) {
     return prefix + result + suffix;
   }
 
+  private String getSetTagForAliasedVariables(String fullName) {
+    int lastDotIdx = fullName.lastIndexOf('.');
+    if (lastDotIdx > 0) {
+      String aliasName = fullName.substring(0, lastDotIdx + 1);
+      Map<String, String> namesToAlias = localContextScope
+        .getCombinedScope()
+        .entrySet()
+        .stream()
+        .filter(entry -> entry.getValue() instanceof DeferredValue)
+        .map(Entry::getKey)
+        .collect(Collectors.toMap(Function.identity(), name -> aliasName + name));
+      return EagerReconstructionUtils.buildSetTag(
+        namesToAlias,
+        JinjavaInterpreter.getCurrent(),
+        false
+      );
+    }
+    return "";
+  }
+
   private boolean differentMacroWithSameNameExists(JinjavaInterpreter interpreter) {
     Context context = interpreter.getContext();
     if (context.getParent() == null) {

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java
@@ -44,7 +44,6 @@
 import java.beans.PropertyDescriptor;
 import java.util.ConcurrentModificationException;
 import java.util.List;
-import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.regex.Matcher;
@@ -151,6 +150,20 @@ public String interpretUnchecked(TagNode tagNode, JinjavaInterpreter interpreter
         String.format("%s in %s", String.join(", ", loopVars), e.getDeferredEvalResult())
       );
     }
+    return renderForCollection(
+      tagNode,
+      interpreter,
+      loopVarsAndExpression.getLeft(),
+      collection
+    );
+  }
+
+  public String renderForCollection(
+    TagNode tagNode,
+    JinjavaInterpreter interpreter,
+    List<String> loopVars,
+    Object collection
+  ) {
     ForLoop loop = ObjectIterator.getLoop(collection);
 
     try (InterpreterScopeClosable c = interpreter.enterScope()) {
@@ -190,8 +203,8 @@ public String interpretUnchecked(TagNode tagNode, JinjavaInterpreter interpreter
         } else {
           for (int loopVarIndex = 0; loopVarIndex < loopVars.size(); loopVarIndex++) {
             String loopVar = loopVars.get(loopVarIndex);
-            if (Map.Entry.class.isAssignableFrom(val.getClass())) {
-              Map.Entry<String, Object> entry = (Entry<String, Object>) val;
+            if (Entry.class.isAssignableFrom(val.getClass())) {
+              Entry<String, Object> entry = (Entry<String, Object>) val;
               Object entryVal = null;
 
               if (loopVars.indexOf(loopVar) == 0) {

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerDoTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerDoTag.java
@@ -40,7 +40,6 @@ public String eagerInterpret(
         EagerContextWatcher
           .EagerChildContextConfig.newBuilder()
           .withTakeNewValue(true)
-          .withCheckForContextChanges(!interpreter.getContext().isDeferredExecutionMode())
           .build()
       );
       PrefixToPreserveState prefixToPreserveState = new PrefixToPreserveState();

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerExecutionResult.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerExecutionResult.java
@@ -4,6 +4,7 @@
 import static com.hubspot.jinjava.util.EagerReconstructionUtils.buildSetTag;
 
 import com.google.common.annotations.Beta;
+import com.hubspot.jinjava.interpret.DeferredLazyReferenceSource;
 import com.hubspot.jinjava.interpret.DeferredValueShadow;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.interpret.LazyReference;
@@ -55,8 +56,13 @@ public PrefixToPreserveState getPrefixToPreserveState() {
       .entrySet()
       .stream()
       .filter(
-        entry ->
-          !(interpreter.getContext().get(entry.getKey()) instanceof DeferredValueShadow)
+        entry -> {
+          Object contextValue = interpreter.getContext().get(entry.getKey());
+          if (contextValue instanceof DeferredLazyReferenceSource) {
+            ((DeferredLazyReferenceSource) contextValue).setReconstructed(true);
+          }
+          return (contextValue != null && !(contextValue instanceof DeferredValueShadow));
+        }
       )
       .collect(Collectors.toList());
     prefixToPreserveState.putAll(

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerForTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerForTag.java
@@ -39,45 +39,63 @@ public EagerForTag(ForTag forTag) {
 
   @Override
   public String innerInterpret(TagNode tagNode, JinjavaInterpreter interpreter) {
-    Set<DeferredToken> addedTokens = new HashSet<>();
-    EagerExecutionResult result = EagerContextWatcher.executeInChildContext(
-      eagerInterpreter -> {
-        EagerExpressionResult expressionResult = EagerExpressionResult.fromSupplier(
-          () -> getTag().interpretUnchecked(tagNode, eagerInterpreter),
-          eagerInterpreter
-        );
-        addedTokens.addAll(eagerInterpreter.getContext().getDeferredTokens());
-        return expressionResult;
-      },
+    Pair<List<String>, String> loopVarsAndExpression = getTag()
+      .getLoopVarsAndExpression((TagToken) tagNode.getMaster());
+    EagerExecutionResult collectionResult = EagerContextWatcher.executeInChildContext(
+      eagerInterpreter ->
+        EagerExpressionResolver.resolveExpression(
+          '[' + loopVarsAndExpression.getRight() + ']',
+          interpreter
+        ),
       interpreter,
       EagerContextWatcher
         .EagerChildContextConfig.newBuilder()
         .withCheckForContextChanges(!interpreter.getContext().isDeferredExecutionMode())
         .build()
     );
-    if (
-      result.getResult().getResolutionState() == ResolutionState.NONE ||
-      (
-        !result.getResult().isFullyResolved() &&
-        !result.getSpeculativeBindings().isEmpty()
-      )
-    ) {
-      EagerReconstructionUtils.resetSpeculativeBindings(interpreter, result);
-      interpreter.getContext().removeDeferredTokens(addedTokens);
-      throw new DeferredValueException(
-        result.getResult().getResolutionState() == ResolutionState.NONE
-          ? result.getResult().toString()
-          : "Modification inside partially evaluated for loop"
-      );
-    }
-    if (result.getResult().isFullyResolved()) {
-      return result.getResult().toString(true);
-    } else {
-      return EagerReconstructionUtils.wrapInChildScope(
-        result.getResult().toString(true),
-        interpreter
+    if (collectionResult.getResult().isFullyResolved()) {
+      Set<DeferredToken> addedTokens = new HashSet<>();
+      EagerExecutionResult result = EagerContextWatcher.executeInChildContext(
+        eagerInterpreter -> {
+          EagerExpressionResult expressionResult = EagerExpressionResult.fromSupplier(
+            () ->
+              getTag()
+                .renderForCollection(
+                  tagNode,
+                  eagerInterpreter,
+                  loopVarsAndExpression.getLeft(),
+                  collectionResult.getResult().toList().get(0)
+                ),
+            eagerInterpreter
+          );
+          addedTokens.addAll(eagerInterpreter.getContext().getDeferredTokens());
+          return expressionResult;
+        },
+        interpreter,
+        EagerContextWatcher.EagerChildContextConfig.newBuilder().build()
       );
+      if (result.getResult().getResolutionState() == ResolutionState.NONE) {
+        EagerReconstructionUtils.resetSpeculativeBindings(interpreter, collectionResult);
+        EagerReconstructionUtils.resetSpeculativeBindings(interpreter, result);
+        interpreter.getContext().removeDeferredTokens(addedTokens);
+        throw new DeferredValueException(result.getResult().toString());
+      }
+      if (result.getResult().isFullyResolved()) {
+        return result.getResult().toString(true);
+      } else {
+        return (
+          result
+            .getPrefixToPreserveState()
+            .withAllInFront(collectionResult.getPrefixToPreserveState()) +
+          EagerReconstructionUtils.wrapInChildScope(
+            result.getResult().toString(true),
+            interpreter
+          )
+        );
+      }
     }
+    EagerReconstructionUtils.resetSpeculativeBindings(interpreter, collectionResult);
+    throw new DeferredValueException(collectionResult.getResult().toString());
   }
 
   @Override

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerImportTag.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerImportTag.java
@@ -191,28 +191,33 @@ private String getFinalOutputWithAlias(
   ) {
     return (
       newPathSetter +
-      getSetTagForDeferredChildBindings(interpreter, currentImportAlias, childBindings) +
       EagerReconstructionUtils.buildSetTag(
         ImmutableMap.of(currentImportAlias, "{}"),
         interpreter,
         true
       ) +
-      wrapInChildScopeIfNecessary(interpreter, output, currentImportAlias) +
+      wrapInChildScope(
+        interpreter,
+        getSetTagForDeferredChildBindings(
+          interpreter,
+          currentImportAlias,
+          childBindings
+        ) +
+        output,
+        currentImportAlias
+      ) +
       initialPathSetter
     );
   }
 
-  private static String wrapInChildScopeIfNecessary(
+  private static String wrapInChildScope(
     JinjavaInterpreter interpreter,
     String output,
     String currentImportAlias
   ) {
     String combined = output + getDoTagToPreserve(interpreter, currentImportAlias);
     // So that any set variables other than the alias won't exist outside the child's scope
-    if (interpreter.getContext().isDeferredExecutionMode()) {
-      return EagerReconstructionUtils.wrapInChildScope(combined, interpreter);
-    }
-    return combined;
+    return EagerReconstructionUtils.wrapInChildScope(combined, interpreter);
   }
 
   private String getSetTagForDeferredChildBindings(
@@ -246,7 +251,11 @@ private String getSetTagForDeferredChildBindings(
       childBindings
         .entrySet()
         .stream()
-        .filter(entry -> entry.getValue() instanceof DeferredValue)
+        .filter(
+          entry ->
+            entry.getValue() instanceof DeferredValue &&
+            ((DeferredValue) entry.getValue()).getOriginalValue() != null
+        )
         .filter(entry -> !interpreter.getContext().containsKey(entry.getKey()))
         .filter(entry -> !entry.getKey().equals(currentImportAlias))
         .collect(

diff --git a/src/main/java/com/hubspot/jinjava/util/EagerReconstructionUtils.java b/src/main/java/com/hubspot/jinjava/util/EagerReconstructionUtils.java
@@ -790,7 +790,8 @@ public static Map<String, String> reconstructDeferredReferences(
    *   * When doing some eager execution and then needing to repeat the same execution in deferred execution mode.
    *   <p>
    *   * When rendering logic which takes place in its own child scope (for tag, macro function, set block) and there
-   *   speculative bindings. These must be deferred and the execution must run again so they don't get reconstructed
+   *   are speculative bindings.
+   *   These must be deferred and the execution must run again, so they don't get reconstructed
    *   within the child scope, and can instead be reconstructed in their original scopes.
    * @param interpreter The JinjavaInterpreter
    * @param eagerExecutionResult The execution result which contains information about which bindings were modified

diff --git a/src/test/java/com/hubspot/jinjava/EagerTest.java b/src/test/java/com/hubspot/jinjava/EagerTest.java
@@ -1145,7 +1145,7 @@ public void itHandlesHigherScopeReferenceModificationSecondPass() {
 
   @Test
   public void itHandlesReferenceModificationWhenSourceIsLost() {
-    expectedTemplateInterpreter.assertExpectedOutput(
+    expectedTemplateInterpreter.assertExpectedOutputNonIdempotent(
       "handles-reference-modification-when-source-is-lost"
     );
   }
@@ -1335,4 +1335,19 @@ public void itDoesNotReconstructVariableInSetInWrongScope() {
       "does-not-reconstruct-variable-in-set-in-wrong-scope"
     );
   }
+
+  @Test
+  public void itRreconstructsValueUsedInDeferredImportedMacro() {
+    expectedTemplateInterpreter.assertExpectedOutputNonIdempotent(
+      "reconstructs-value-used-in-deferred-imported-macro"
+    );
+  }
+
+  @Test
+  public void itRreconstructsValueUsedInDeferredImportedMacroSecondPass() {
+    interpreter.getContext().put("deferred", "resolved");
+    expectedTemplateInterpreter.assertExpectedOutput(
+      "reconstructs-value-used-in-deferred-imported-macro.expected"
+    );
+  }
 }
diff --git a/src/test/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunctionTest.java b/src/test/java/com/hubspot/jinjava/lib/fn/eager/EagerMacroFunctionTest.java
@@ -65,6 +65,7 @@ public void itResolvesFromContext() {
 
   @Test
   public void itReconstructsForAliasedName() {
+    context.remove("deferred");
     String name = "foo";
     String fullName = "local." + name;
     String codeFormat = "{%% macro %s(bar) %%}It's: {{ bar }}{%% endmacro %%}";