Merge pull request #692 from HubSpot/separate-if-branches

[Eager Execution] Execute deferred branches of if tag in separate contexts

Author: jasmith-hs

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

@@ -184,6 +184,9 @@ private void setVariable(JinjavaInterpreter interpreter, String var, Object valu
         ((Namespace) namespace).put(varArray[1], value);
         return;
       }
+      if (namespace instanceof DeferredValue) {
+        throw new DeferredValueException("Deferred Namespace");
+      }
     }
     interpreter.getContext().put(var, value);
   }

src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerIfTag.java

@@ -1,5 +1,6 @@
 package com.hubspot.jinjava.lib.tag.eager;
 
+import com.hubspot.jinjava.interpret.DeferredValue;
 import com.hubspot.jinjava.interpret.DeferredValueException;
 import com.hubspot.jinjava.interpret.InterpretException;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
@@ -98,51 +99,113 @@ public String eagerRenderBranches(
     boolean definitelyExecuted = false;
     StringBuilder sb = new StringBuilder();
     sb.append(getEagerImage(buildToken(tagNode, e, deferredLineNumber), interpreter));
+    int branchStart = 0;
+    int childrenSize = tagNode.getChildren().size();
+    while (branchStart < childrenSize) {
+      int branchEnd = findNextElseToken(tagNode, branchStart);
+      if (!definitelyDrop) {
+        int finalBranchStart = branchStart;
+        EagerExecutionResult result = executeInChildContext(
+          eagerInterpreter ->
+            EagerExpressionResult.fromString(
+              evaluateBranch(tagNode, finalBranchStart, branchEnd, interpreter)
+            ),
+          interpreter,
+          false,
+          false,
+          true
+        );
+        sb.append(result.getResult());
+        resetBindingsForNextBranch(interpreter, result);
+      }
+      if (branchEnd >= childrenSize || definitelyExecuted) {
+        break;
+      }
+      TagNode caseNode = (TagNode) tagNode.getChildren().get(branchEnd);
+      definitelyDrop =
+        caseNode.getName().equals(ElseIfTag.TAG_NAME) &&
+        shouldDropBranch(caseNode, interpreter, deferredLineNumber);
+      if (!definitelyDrop) {
+        definitelyExecuted =
+          caseNode.getName().equals(ElseTag.TAG_NAME) ||
+          isDefinitelyExecuted(caseNode, interpreter, deferredLineNumber);
+        if (definitelyExecuted) {
+          sb.append(
+            String.format(
+              "%s else %s",
+              caseNode.getSymbols().getExpressionStartWithTag(),
+              caseNode.getSymbols().getExpressionEndWithTag()
+            )
+          );
+        } else {
+          sb.append(
+            getEagerImage(buildToken(caseNode, e, deferredLineNumber), interpreter)
+          );
+        }
+      }
+      branchStart = branchEnd + 1;
+    }
+    return sb.toString();
+  }
 
-    for (Node child : tagNode.getChildren()) {
-      if (TagNode.class.isAssignableFrom(child.getClass())) {
-        TagNode childTagNode = (TagNode) child;
-        if (
-          childTagNode.getName().equals(ElseIfTag.TAG_NAME) ||
-          childTagNode.getName().equals(ElseTag.TAG_NAME)
-        ) {
-          if (definitelyExecuted) {
-            break;
-          }
-          definitelyDrop =
-            childTagNode.getName().equals(ElseIfTag.TAG_NAME) &&
-            shouldDropBranch(childTagNode, interpreter, deferredLineNumber);
-          if (!definitelyDrop) {
-            definitelyExecuted =
-              childTagNode.getName().equals(ElseTag.TAG_NAME) ||
-              isDefinitelyExecuted(childTagNode, interpreter, deferredLineNumber);
-            if (definitelyExecuted) {
-              sb.append(
-                String.format(
-                  "%s else %s",
-                  childTagNode.getSymbols().getExpressionStartWithTag(),
-                  childTagNode.getSymbols().getExpressionEndWithTag()
-                )
-              );
-            } else {
-              sb.append(
-                getEagerImage(
-                  buildToken(childTagNode, e, deferredLineNumber),
-                  interpreter
-                )
+  private void resetBindingsForNextBranch(
+    JinjavaInterpreter interpreter,
+    EagerExecutionResult result
+  ) {
+    result
+      .getSpeculativeBindings()
+      .keySet()
+      .stream()
+      .filter(
+        key ->
+          interpreter.getContext().containsKey(key) &&
+          interpreter.getContext().get(key) instanceof DeferredValue
+      )
+      .forEach(
+        key -> {
+          if (
+            ((DeferredValue) interpreter.getContext().get(key)).getOriginalValue() != null
+          ) {
+            interpreter
+              .getContext()
+              .put(
+                key,
+                ((DeferredValue) interpreter.getContext().get(key)).getOriginalValue()
               );
-            }
           }
-          continue;
         }
-      }
-      if (!definitelyDrop) {
-        sb.append(child.render(interpreter).getValue());
-      }
+      );
+  }
+
+  private String evaluateBranch(
+    TagNode tagNode,
+    int startIdx,
+    int endIdx,
+    JinjavaInterpreter interpreter
+  ) {
+    StringBuilder sb = new StringBuilder();
+    for (int i = startIdx; i < endIdx; i++) {
+      Node child = tagNode.getChildren().get(i);
+      sb.append(child.render(interpreter).getValue());
     }
     return sb.toString();
   }
 
+  private int findNextElseToken(TagNode tagNode, int startIdx) {
+    int i;
+    for (i = startIdx; i < tagNode.getChildren().size(); i++) {
+      Node childNode = tagNode.getChildren().get(i);
+      if (
+        (TagNode.class.isAssignableFrom(childNode.getClass())) &&
+        childNode.getName().equals(ElseIfTag.TAG_NAME) ||
+        childNode.getName().equals(ElseTag.TAG_NAME)
+      ) {
+        return i;
+      }
+    }
+    return i;
+  }
+
   private boolean shouldDropBranch(
     TagNode tagNode,
     JinjavaInterpreter eagerInterpreter,

src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerSetTag.java

@@ -133,6 +133,21 @@ public String getEagerTagImage(TagToken tagToken, JinjavaInterpreter interpreter
         )
       );
     }
+    if (
+      eagerExecutionResult.getResult().isFullyResolved() &&
+      interpreter.getContext().isDeferredExecutionMode()
+    ) {
+      try {
+        getTag()
+          .executeSet(
+            tagToken,
+            interpreter,
+            varTokens,
+            eagerExecutionResult.getResult().toList(),
+            true
+          );
+      } catch (DeferredValueException ignored) {}
+    }
     return wrapInAutoEscapeIfNeeded(
       prefixToPreserveState + joiner.toString() + suffixToPreserveState.toString(),
       interpreter

src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerTagDecorator.java

@@ -267,7 +267,13 @@ public static EagerExecutionResult executeInChildContext(
         .entrySet()
         .stream()
         .filter(
-          entry -> !entry.getValue().equals(interpreter.getContext().get(entry.getKey()))
+          entry ->
+            entry.getValue() != null &&
+            !entry.getValue().equals(interpreter.getContext().get(entry.getKey()))
+        )
+        .filter(
+          entry ->
+            !(interpreter.getContext().get(entry.getKey()) instanceof DeferredValue)
         )
         .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
     if (checkForContextChanges) {

src/test/java/com/hubspot/jinjava/EagerTest.java

@@ -423,7 +423,7 @@ public void itMarksVariablesSetInDeferredBlockAsDeferred() {
     assertThat(localContext).containsKey("varSetInside");
     Object varSetInside = localContext.get("varSetInside");
     assertThat(varSetInside).isInstanceOf(DeferredValue.class);
-    assertThat(output).contains("{{ varSetInside }}");
+    assertThat(output).contains("{{ varSetInside }}").contains("xyz");
     assertThat(context.get("a")).isInstanceOf(DeferredValue.class);
     assertThat(context.get("b")).isInstanceOf(DeferredValue.class);
     assertThat(context.get("c")).isInstanceOf(DeferredValue.class);

src/test/resources/deferred/set-in-deferred.jinja

@@ -1,6 +1,7 @@
 {% set reference = deferredValue %}
 {% if reference == 'resolved' %}
    {% set varSetInside = 'set inside' %}
-   {% set a, b, c = ['x','y','z'] %}
+   {% set a, b, c = 'x','y','z' %}
+   {{ a ~ b ~ c }}
 {% endif %}
 {{ varSetInside }}

src/test/resources/eager/defers-eager-child-scoped-vars.expected.jinja

@@ -1,5 +1,5 @@
 {% if deferred %}
 {% set foo = 1 %}
-{{ foo }}
+1
 {% endif %}
 We don't know if someone wanted to access {{ foo }} here!

src/test/resources/eager/defers-on-immutable-mode.expected.jinja

@@ -1,7 +1,7 @@
 {% set foo = 1 %}{% if deferred %}
 {% set foo = 2 %}
 {% else %}
-{% set foo = foo + 2 %}//best we can do
+{% set foo = 3 %}
 {% endif %}
 {{ foo }}
 

src/test/resources/eager/defers-on-immutable-mode.jinja

@@ -2,7 +2,7 @@
 {% if deferred %}
 {% set foo = foo + 1 %}
 {% else %}
-{% set foo = foo + 2 %}//best we can do
+{% set foo = foo + 2 %}
 {% endif %}
 {{ foo }}