Defer variables used in deferred

src/main/java/com/hubspot/jinjava/interpret/Context.java

@@ -30,6 +30,7 @@
 import com.hubspot.jinjava.lib.tag.Tag;
 import com.hubspot.jinjava.lib.tag.TagLibrary;
 import com.hubspot.jinjava.tree.Node;
+import com.hubspot.jinjava.util.DeferredValueUtils;
 import com.hubspot.jinjava.util.ScopeMap;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -275,10 +276,20 @@ public void addResolvedFunction(String function) {
     }
   }
 
-  public void addDeferredNode(Node node) {
+  public void handleDeferredNode(Node node) {
     deferredNodes.add(node);
+    Set<String> deferredProps = DeferredValueUtils.findAndMarkDeferredProperties(this);
     if (getParent() != null) {
-      getParent().addDeferredNode(node);
+      Context parent = getParent();
+      //Ignore global context
+      if (parent.getParent() != null) {
+        //Place deferred values on the parent context
+        deferredProps
+          .stream()
+          .filter(key -> !parent.containsKey(key))
+          .forEach(key -> parent.put(key, this.get(key)));
+        getParent().handleDeferredNode(node);
+      }
     }
   }
 

src/main/java/com/hubspot/jinjava/interpret/JinjavaInterpreter.java

@@ -267,7 +267,7 @@ public String render(Node root, boolean processExtendRoots) {
         try {
           out = node.render(this);
         } catch (DeferredValueException e) {
-          context.addDeferredNode(node);
+          context.handleDeferredNode(node);
           out = new RenderedOutputNode(node.getMaster().getImage());
         }
         context.popRenderStack();

src/main/java/com/hubspot/jinjava/lib/tag/ImportTag.java

@@ -150,7 +150,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
         node
           .getChildren()
           .forEach(
-            deferredChild -> interpreter.getContext().addDeferredNode(deferredChild)
+            deferredChild -> interpreter.getContext().handleDeferredNode(deferredChild)
           );
         if (StringUtils.isBlank(contextVar)) {
           for (MacroFunction macro : child.getContext().getGlobalMacros().values()) {
@@ -173,10 +173,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
           }
           childBindings.remove(Context.GLOBAL_MACROS_SCOPE_KEY);
           childBindings.remove(Context.IMPORT_RESOURCE_PATH_KEY);
-          for (String key : childBindings.keySet()) {
-            childBindings.put(key, DeferredValue.instance());
-          }
-          interpreter.getContext().put(contextVar, childBindings);
+          interpreter.getContext().put(contextVar, DeferredValue.instance(childBindings));
         }
 
         throw new DeferredValueException(

src/main/java/com/hubspot/jinjava/lib/tag/SetTag.java

@@ -126,10 +126,20 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
 
         for (int i = 0; i < varTokens.length; i++) {
           String varItem = varTokens[i].trim();
+          if (interpreter.getContext().containsKey(varItem)) {
+            if (interpreter.getContext().get(varItem) instanceof DeferredValue) {
+              throw new DeferredValueException(varItem);
+            }
+          }
           interpreter.getContext().put(varItem, exprVals.get(i));
         }
       } else {
         // handle single variable assignment
+        if (interpreter.getContext().containsKey(var)) {
+          if (interpreter.getContext().get(var) instanceof DeferredValue) {
+            throw new DeferredValueException(var);
+          }
+        }
         interpreter
           .getContext()
           .put(var, interpreter.resolveELExpression(expr, tagNode.getLineNumber()));

src/main/java/com/hubspot/jinjava/tree/ExpressionNode.java

@@ -42,7 +42,7 @@ public OutputNode render(JinjavaInterpreter interpreter) {
     try {
       var = interpreter.resolveELExpression(master.getExpr(), getLineNumber());
     } catch (DeferredValueException e) {
-      interpreter.getContext().addDeferredNode(this);
+      interpreter.getContext().handleDeferredNode(this);
       var = master.getImage();
     }
 

src/main/java/com/hubspot/jinjava/tree/TagNode.java

@@ -51,7 +51,7 @@ public OutputNode render(JinjavaInterpreter interpreter) {
     try {
       return tag.interpretOutput(this, interpreter);
     } catch (DeferredValueException e) {
-      interpreter.getContext().addDeferredNode(this);
+      interpreter.getContext().handleDeferredNode(this);
       return new RenderedOutputNode(reconstructImage());
     } catch (InterpretException | InvalidInputException | InvalidArgumentException e) {
       throw e;

src/main/java/com/hubspot/jinjava/util/DeferredValueUtils.java

@@ -0,0 +1,207 @@
+package com.hubspot.jinjava.util;
+
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
+import com.hubspot.jinjava.interpret.Context;
+import com.hubspot.jinjava.interpret.DeferredValue;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.lib.tag.SetTag;
+import com.hubspot.jinjava.tree.ExpressionNode;
+import com.hubspot.jinjava.tree.Node;
+import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.tree.TextNode;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.StringJoiner;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+public class DeferredValueUtils {
+  private static final String TEMPLATE_TAG_REGEX = "(\\w+(?:\\.\\w+)*)";
+  private static final Pattern TEMPLATE_TAG_PATTERN = Pattern.compile(TEMPLATE_TAG_REGEX);
+
+  public static HashMap<String, Object> getDeferredContextWithOriginalValues(
+    Context context
+  ) {
+    return getDeferredContextWithOriginalValues(context, ImmutableSet.of());
+  }
+
+  //The context needed for a second render
+  //Ignores deferred properties with no originalValue
+  //Optionally only keep keys in keysToKeep
+  public static HashMap<String, Object> getDeferredContextWithOriginalValues(
+    Context context,
+    Set<String> keysToKeep
+  ) {
+    HashMap<String, Object> deferredContext = new HashMap<>(context.size());
+    context.forEach(
+      (contextKey, contextItem) -> {
+        if (keysToKeep.size() > 0 && !keysToKeep.contains(contextKey)) {
+          return;
+        }
+        if (contextItem instanceof DeferredValue) {
+          if (((DeferredValue) contextItem).getOriginalValue() != null) {
+            deferredContext.put(
+              contextKey,
+              ((DeferredValue) contextItem).getOriginalValue()
+            );
+          }
+        }
+      }
+    );
+    return deferredContext;
+  }
+
+  public static Set<String> findAndMarkDeferredProperties(Context context) {
+    String templateSource = rebuildTemplateForNodes(context.getDeferredNodes());
+    Set<String> deferredProps = getPropertiesUsedInDeferredNodes(context, templateSource);
+    Set<String> setProps = getPropertiesSetInDeferredNodes(templateSource);
+
+    markDeferredProperties(context, Sets.union(deferredProps, setProps));
+
+    return deferredProps;
+  }
+
+  public static Set<String> getPropertiesSetInDeferredNodes(String templateSource) {
+    return findSetProperties(templateSource);
+  }
+
+  public static Set<DeferredTag> getDeferredTags(Set<Node> deferredNodes) {
+    return getDeferredTags(new LinkedList<>(deferredNodes), 0);
+  }
+
+  public static Set<String> getPropertiesUsedInDeferredNodes(
+    Context context,
+    String templateSource
+  ) {
+    Set<String> propertiesUsed = findUsedProperties(templateSource);
+    return propertiesUsed
+      .stream()
+      .map(prop -> prop.split("\\.", 2)[0]) // split accesses on .prop
+      .filter(context::containsKey)
+      .collect(Collectors.toSet());
+  }
+
+  private static void markDeferredProperties(Context context, Set<String> props) {
+    props
+      .stream()
+      .filter(prop -> !(context.get(prop) instanceof DeferredValue))
+      .forEach(
+        prop -> {
+          if (context.get(prop) != null) {
+            context.put(prop, DeferredValue.instance(context.get(prop)));
+          } else {
+            //Handle set props
+            context.put(prop, DeferredValue.instance());
+          }
+        }
+      );
+  }
+
+  private static Set<DeferredTag> getDeferredTags(List<Node> nodes, int depth) {
+    // precaution - templates are parsed with this render depth so in theory the depth should never be exceeded
+    Set<DeferredTag> deferredTags = new HashSet<>();
+    int maxRenderDepth = JinjavaInterpreter.getCurrent() == null
+      ? 3
+      : JinjavaInterpreter.getCurrent().getConfig().getMaxRenderDepth();
+    if (depth > maxRenderDepth) {
+      return deferredTags;
+    }
+    for (Node node : nodes) {
+      getDeferredTags(node).ifPresent(deferredTags::addAll);
+      deferredTags.addAll(getDeferredTags(node.getChildren(), depth + 1));
+    }
+    return deferredTags;
+  }
+
+  private static String rebuildTemplateForNodes(Set<Node> nodes) {
+    StringJoiner joiner = new StringJoiner(" ");
+    getDeferredTags(nodes).stream().map(DeferredTag::getTag).forEach(joiner::add);
+    return joiner.toString();
+  }
+
+  private static Set<String> findUsedProperties(String templateSource) {
+    Matcher matcher = TEMPLATE_TAG_PATTERN.matcher(templateSource);
+    Set<String> tags = Sets.newHashSet();
+    while (matcher.find()) {
+      tags.add(matcher.group(1));
+    }
+    return tags;
+  }
+
+  private static Set<String> findSetProperties(String templateSource) {
+    Set<String> tags = Sets.newHashSet();
+    String[] lines = templateSource.split("\n");
+    for (String line : lines) {
+      line = line.trim();
+      if (line.contains(SetTag.TAG_NAME + " ")) {
+        tags.addAll(findUsedProperties(line));
+      }
+    }
+
+    return tags;
+  }
+
+  private static Optional<Set<DeferredTag>> getDeferredTags(Node deferredNode) {
+    if (deferredNode instanceof TextNode || deferredNode.getMaster() == null) {
+      return Optional.empty();
+    }
+
+    String nodeImage = deferredNode.getMaster().getImage();
+    if (Strings.nullToEmpty(nodeImage).trim().isEmpty()) {
+      return Optional.empty();
+    }
+
+    Set<DeferredTag> deferredTags = new HashSet<>();
+    deferredTags.add(
+      new DeferredTag().setTag(nodeImage).setNormalizedTag(getNormalizedTag(deferredNode))
+    );
+
+    if (deferredNode instanceof TagNode) {
+      TagNode tagNode = (TagNode) deferredNode;
+      if (tagNode.getEndName() != null) {
+        String endTag = tagNode.reconstructEnd();
+        deferredTags.add(new DeferredTag().setTag(endTag).setNormalizedTag(endTag));
+      }
+    }
+
+    return Optional.of(deferredTags);
+  }
+
+  private static String getNormalizedTag(Node node) {
+    if (node instanceof ExpressionNode) {
+      return node.toString().replaceAll("\\s+", "");
+    }
+
+    return node.getMaster().getImage();
+  }
+
+  private static class DeferredTag {
+    String tag;
+    String normalizedTag;
+
+    public String getTag() {
+      return tag;
+    }
+
+    public DeferredTag setTag(String tag) {
+      this.tag = tag;
+      return this;
+    }
+
+    public String getNormalizedTag() {
+      return normalizedTag;
+    }
+
+    public DeferredTag setNormalizedTag(String normalizedTag) {
+      this.normalizedTag = normalizedTag;
+      return this;
+    }
+  }
+}