From 1f88ba3f6bcde69ee1270f09c4cbfd79dcf1ac22 Mon Sep 17 00:00:00 2001
From: Jeff Boulter <jboulter@hubspot.com>
Date: Mon, 15 Apr 2024 17:36:40 -0400
Subject: [PATCH] find first valid deferred words

---
 .../jinjava/lib/tag/eager/DeferredToken.java  | 20 +++++++++++++++---
 .../java/com/hubspot/jinjava/EagerTest.java   | 21 +++++++++++++++++++
 .../jinjava/util/DeferredValueUtilsTest.java  | 16 ++++++++++++++
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/DeferredToken.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/DeferredToken.java
index 963f1e346..1abf83c6f 100644
--- a/src/main/java/com/hubspot/jinjava/lib/tag/eager/DeferredToken.java
+++ b/src/main/java/com/hubspot/jinjava/lib/tag/eager/DeferredToken.java
@@ -1,6 +1,7 @@
 package com.hubspot.jinjava.lib.tag.eager;
 
 import com.google.common.annotations.Beta;
+import com.google.common.base.Strings;
 import com.hubspot.jinjava.interpret.CallStack;
 import com.hubspot.jinjava.interpret.Context;
 import com.hubspot.jinjava.interpret.DeferredLazyReference;
@@ -13,9 +14,11 @@
 import com.hubspot.jinjava.tree.parse.TokenScannerSymbols;
 import com.hubspot.jinjava.util.EagerExpressionResolver;
 import java.lang.reflect.InvocationTargetException;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -40,7 +43,8 @@ public DeferredToken build() {
         token,
         usedDeferredWords != null
           ? usedDeferredWords
-            .map(prop -> prop.split("\\.", 2)[0])
+            .map(DeferredToken::splitToken)
+            .map(DeferredToken::getFirstNonEmptyToken)
             .distinct()
             .filter(word ->
               interpreter == null ||
@@ -50,7 +54,8 @@ public DeferredToken build() {
           : Collections.emptySet(),
         setDeferredWords != null
           ? setDeferredWords
-            .map(prop -> prop.split("\\.", 2)[0])
+            .map(DeferredToken::splitToken)
+            .map(DeferredToken::getFirstNonEmptyToken)
             .collect(Collectors.toSet())
           : Collections.emptySet(),
         acquireImportResourcePath(),
@@ -413,10 +418,19 @@ private static CallStack acquireMacroStack() {
       .orElse(null);
   }
 
+  private static String getFirstNonEmptyToken(List<String> strings) {
+    return Strings.isNullOrEmpty(strings.get(0)) ? strings.get(1) : strings.get(0);
+  }
+
+  public static List<String> splitToken(String token) {
+    return Arrays.asList(token.split("\\.", 2));
+  }
+
   public static Set<String> getBases(Set<String> original) {
     return original
       .stream()
-      .map(prop -> prop.split("\\.", 2)[0])
+      .map(DeferredToken::splitToken)
+      .map(prop -> prop.get(0))
       .collect(Collectors.toSet());
   }
 }
diff --git a/src/test/java/com/hubspot/jinjava/EagerTest.java b/src/test/java/com/hubspot/jinjava/EagerTest.java
index 4dd3010a7..e59f6b32a 100644
--- a/src/test/java/com/hubspot/jinjava/EagerTest.java
+++ b/src/test/java/com/hubspot/jinjava/EagerTest.java
@@ -517,6 +517,27 @@ public void itEvaluatesNonEagerSet() {
       .contains("deferred");
   }
 
+  @Test
+  public void itDefersArrayAccess() {
+    expectedTemplateInterpreter.assertExpectedOutput("evaluates-non-eager-set");
+    assertThat(
+      localContext
+        .getDeferredTokens()
+        .stream()
+        .flatMap(deferredToken -> deferredToken.getSetDeferredWords().stream())
+        .collect(Collectors.toSet())
+    )
+      .isEmpty();
+    assertThat(
+      localContext
+        .getDeferredTokens()
+        .stream()
+        .flatMap(deferredToken -> deferredToken.getUsedDeferredWords().stream())
+        .collect(Collectors.toSet())
+    )
+      .contains("deferred");
+  }
+
   @Test
   public void itDefersOnImmutableMode() {
     expectedTemplateInterpreter.assertExpectedOutput("defers-on-immutable-mode");
diff --git a/src/test/java/com/hubspot/jinjava/util/DeferredValueUtilsTest.java b/src/test/java/com/hubspot/jinjava/util/DeferredValueUtilsTest.java
index f232ed034..f037e4c90 100644
--- a/src/test/java/com/hubspot/jinjava/util/DeferredValueUtilsTest.java
+++ b/src/test/java/com/hubspot/jinjava/util/DeferredValueUtilsTest.java
@@ -246,6 +246,22 @@ public void itDefersUsedWordsInDeferredTokens() {
     assertThat(context.containsKey("int")).isFalse();
   }
 
+  @Test
+  public void itFindsFirstValidDeferredWords() {
+    DeferredToken deferredToken = DeferredToken
+      .builderFromToken(
+        new ExpressionToken("{{ blah }}", 1, 1, new DefaultTokenScannerSymbols())
+      )
+      .addUsedDeferredWords(ImmutableSet.of("deferred", ".attribute1"))
+      .addSetDeferredWords(ImmutableSet.of("deferred", ".attribute2"))
+      .build();
+
+    assertThat(deferredToken.getUsedDeferredWords())
+      .isEqualTo(ImmutableSet.of("deferred", "attribute1"));
+    assertThat(deferredToken.getSetDeferredWords())
+      .isEqualTo(ImmutableSet.of("deferred", "attribute2"));
+  }
+
   private Context getContext(List<? extends Node> nodes) {
     return getContext(nodes, Optional.empty());
   }