Retain both the original deferred words and their bases, and expose them

with separate methods

Author: jasmith-hs

src/main/java/com/hubspot/jinjava/lib/fn/MacroFunction.java

@@ -237,11 +237,8 @@ private boolean alreadyDeferredInEarlierCall(
           Objects.equals(importResourcePath, deferredToken.getImportResourcePath())
         )
         .anyMatch(deferredToken ->
-          deferredToken.getSetDeferredWords().contains(key) ||
-          deferredToken
-            .getUsedDeferredWords()
-            .stream()
-            .anyMatch(used -> key.equals(used.split("\\.", 2)[0]))
+          deferredToken.getSetDeferredBases().contains(key) ||
+          deferredToken.getUsedDeferredBases().contains(key)
         );
     }
     return false;

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

@@ -2,6 +2,7 @@
 
 import com.google.common.annotations.Beta;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableSet;
 import com.hubspot.jinjava.interpret.CallStack;
 import com.hubspot.jinjava.interpret.Context;
 import com.hubspot.jinjava.interpret.DeferredLazyReference;
@@ -31,34 +32,18 @@ public class DeferredToken {
   public static class DeferredTokenBuilder {
 
     private final Token token;
-    private Stream<String> usedDeferredWords;
-    private Stream<String> setDeferredWords;
+    private ImmutableSet.Builder<String> usedDeferredWords;
+    private ImmutableSet.Builder<String> setDeferredWords;
 
     private DeferredTokenBuilder(Token token) {
       this.token = token;
     }
 
     public DeferredToken build() {
-      JinjavaInterpreter interpreter = JinjavaInterpreter.getCurrent();
       return new DeferredToken(
         token,
-        usedDeferredWords != null
-          ? usedDeferredWords
-            .map(DeferredToken::splitToken)
-            .map(DeferredToken::getFirstNonEmptyToken)
-            .distinct()
-            .filter(word ->
-              interpreter == null ||
-              !(interpreter.getContext().get(word) instanceof DeferredMacroValueImpl)
-            )
-            .collect(Collectors.toSet())
-          : Collections.emptySet(),
-        setDeferredWords != null
-          ? setDeferredWords
-            .map(DeferredToken::splitToken)
-            .map(DeferredToken::getFirstNonEmptyToken)
-            .collect(Collectors.toSet())
-          : Collections.emptySet(),
+        usedDeferredWords != null ? usedDeferredWords.build() : Collections.emptySet(),
+        setDeferredWords != null ? setDeferredWords.build() : Collections.emptySet(),
         acquireImportResourcePath(),
         acquireMacroStack()
       );
@@ -67,34 +52,40 @@ public DeferredToken build() {
     public DeferredTokenBuilder addUsedDeferredWords(
       Collection<String> usedDeferredWordsToAdd
     ) {
-      return addUsedDeferredWords(usedDeferredWordsToAdd.stream());
+      if (usedDeferredWords == null) {
+        usedDeferredWords = ImmutableSet.builder();
+      }
+      usedDeferredWords.addAll(usedDeferredWordsToAdd);
+      return this;
     }
 
     public DeferredTokenBuilder addUsedDeferredWords(
       Stream<String> usedDeferredWordsToAdd
     ) {
       if (usedDeferredWords == null) {
-        usedDeferredWords = usedDeferredWordsToAdd;
-      } else {
-        usedDeferredWords = Stream.concat(usedDeferredWords, usedDeferredWordsToAdd);
+        usedDeferredWords = ImmutableSet.builder();
       }
+      usedDeferredWordsToAdd.forEach(usedDeferredWords::add);
       return this;
     }
 
     public DeferredTokenBuilder addSetDeferredWords(
       Collection<String> setDeferredWordsToAdd
     ) {
-      return addSetDeferredWords(setDeferredWordsToAdd.stream());
+      if (setDeferredWords == null) {
+        setDeferredWords = ImmutableSet.builder();
+      }
+      setDeferredWords.addAll(setDeferredWordsToAdd);
+      return this;
     }
 
     public DeferredTokenBuilder addSetDeferredWords(
       Stream<String> setDeferredWordsToAdd
     ) {
       if (setDeferredWords == null) {
-        setDeferredWords = setDeferredWordsToAdd;
-      } else {
-        setDeferredWords = Stream.concat(setDeferredWords, setDeferredWordsToAdd);
+        setDeferredWords = ImmutableSet.builder();
       }
+      setDeferredWordsToAdd.forEach(setDeferredWords::add);
       return this;
     }
   }
@@ -103,9 +94,10 @@ public DeferredTokenBuilder addSetDeferredWords(
   // These words aren't yet DeferredValues, but are unresolved
   // so they should be replaced with DeferredValueImpls if they exist in the context
   private final Set<String> usedDeferredWords;
-
+  private final Set<String> usedDeferredBases;
   // These words are those which will be set to a value which has been deferred.
   private final Set<String> setDeferredWords;
+  private final Set<String> setDeferredBases;
 
   // Used to determine the combine scope
   private final CallStack macroStack;
@@ -211,23 +203,45 @@ public DeferredToken(
   ) {
     this(
       token,
-      getBases(usedDeferredWords),
-      getBases(setDeferredWords),
+      usedDeferredWords,
+      setDeferredWords,
       acquireImportResourcePath(),
       acquireMacroStack()
     );
   }
 
   private DeferredToken(
     Token token,
-    Set<String> usedDeferredWordBases,
-    Set<String> setDeferredWordBases,
+    Set<String> usedDeferredWords,
+    Set<String> setDeferredWords,
     String importResourcePath,
     CallStack macroStack
   ) {
+    JinjavaInterpreter interpreter = JinjavaInterpreter.getCurrent();
     this.token = token;
-    this.usedDeferredWords = usedDeferredWordBases;
-    this.setDeferredWords = setDeferredWordBases;
+    this.usedDeferredBases =
+      usedDeferredWords.isEmpty()
+        ? Collections.emptySet()
+        : usedDeferredWords
+          .stream()
+          .map(DeferredToken::splitToken)
+          .map(DeferredToken::getFirstNonEmptyToken)
+          .distinct()
+          .filter(word ->
+            interpreter == null ||
+            !(interpreter.getContext().get(word) instanceof DeferredMacroValueImpl)
+          )
+          .collect(Collectors.toSet());
+    this.usedDeferredWords = usedDeferredWords;
+    this.setDeferredBases =
+      setDeferredWords.isEmpty()
+        ? Collections.emptySet()
+        : setDeferredWords
+          .stream()
+          .map(DeferredToken::splitToken)
+          .map(DeferredToken::getFirstNonEmptyToken)
+          .collect(Collectors.toSet());
+    this.setDeferredWords = setDeferredWords;
     this.importResourcePath = importResourcePath;
     this.macroStack = macroStack;
   }
@@ -240,10 +254,18 @@ public Set<String> getUsedDeferredWords() {
     return usedDeferredWords;
   }
 
+  public Set<String> getUsedDeferredBases() {
+    return usedDeferredBases;
+  }
+
   public Set<String> getSetDeferredWords() {
     return setDeferredWords;
   }
 
+  public Set<String> getSetDeferredBases() {
+    return setDeferredBases;
+  }
+
   public String getImportResourcePath() {
     return importResourcePath;
   }
@@ -255,7 +277,7 @@ public CallStack getMacroStack() {
   public void addTo(Context context) {
     addTo(
       context,
-      usedDeferredWords
+      usedDeferredBases
         .stream()
         .filter(word -> {
           Object value = context.get(word);
@@ -285,7 +307,7 @@ private void deferPropertiesOnContext(
   ) {
     if (isInSameScope(context)) {
       // set props are only deferred when within the scope which the variable is set in
-      markDeferredWordsAndFindSources(context, getSetDeferredWords(), true);
+      markDeferredWordsAndFindSources(context, getSetDeferredBases(), true);
     }
     wordsWithoutDeferredSource.forEach(word -> deferDuplicatePointers(context, word));
     wordsWithoutDeferredSource.removeAll(
@@ -424,12 +446,4 @@ private static String getFirstNonEmptyToken(List<String> strings) {
   public static List<String> splitToken(String token) {
     return Arrays.asList(token.split("\\."));
   }
-
-  public static Set<String> getBases(Set<String> original) {
-    return original
-      .stream()
-      .map(DeferredToken::splitToken)
-      .map(prop -> prop.get(0))
-      .collect(Collectors.toSet());
-  }
 }

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

@@ -795,13 +795,13 @@ public static Map<String, String> handleDeferredTokenAndReconstructReferences(
     deferredToken.addTo(interpreter.getContext());
     return reconstructDeferredReferences(
       interpreter,
-      deferredToken.getUsedDeferredWords()
+      deferredToken.getUsedDeferredBases()
     );
   }
 
   public static Map<String, String> reconstructDeferredReferences(
     JinjavaInterpreter interpreter,
-    Set<String> usedDeferredWords
+    Set<String> usedDeferredBases
   ) {
     return interpreter
       .getContext()
@@ -815,7 +815,7 @@ public static Map<String, String> reconstructDeferredReferences(
       .filter(entry ->
         // Always reconstruct the DeferredLazyReferenceSource, but only reconstruct DeferredLazyReferences when they are used
         entry.getValue() instanceof DeferredLazyReferenceSource ||
-        usedDeferredWords.contains(entry.getKey())
+        usedDeferredBases.contains(entry.getKey())
       )
       .peek(entry -> ((OneTimeReconstructible) entry.getValue()).setReconstructed(true))
       .map(entry ->

src/test/java/com/hubspot/jinjava/lib/tag/eager/EagerSetTagTest.java

@@ -159,6 +159,14 @@ public void itDefersBlockWithFilter() {
         .stream()
         .flatMap(deferredToken -> deferredToken.getUsedDeferredWords().stream())
         .collect(Collectors.toSet())
+    )
+      .containsExactlyInAnyOrder("deferred", "foo", "add.filter");
+    assertThat(
+      context
+        .getDeferredTokens()
+        .stream()
+        .flatMap(deferredToken -> deferredToken.getUsedDeferredBases().stream())
+        .collect(Collectors.toSet())
     )
       .containsExactlyInAnyOrder("deferred", "foo", "add");
   }
@@ -211,6 +219,14 @@ public void itDefersInDeferredExecutionModeWithFilter() {
         .stream()
         .flatMap(deferredToken -> deferredToken.getUsedDeferredWords().stream())
         .collect(Collectors.toSet())
+    )
+      .containsExactlyInAnyOrder("deferred", "foo", "add.filter");
+    assertThat(
+      context
+        .getDeferredTokens()
+        .stream()
+        .flatMap(deferredToken -> deferredToken.getUsedDeferredBases().stream())
+        .collect(Collectors.toSet())
     )
       .containsExactlyInAnyOrder("deferred", "foo", "add");
     context.remove("foo");

src/test/java/com/hubspot/jinjava/util/DeferredValueUtilsTest.java

@@ -256,9 +256,9 @@ public void itFindsFirstValidDeferredWords() {
       .addSetDeferredWords(ImmutableSet.of("deferred", ".attribute2"))
       .build();
 
-    assertThat(deferredToken.getUsedDeferredWords())
+    assertThat(deferredToken.getUsedDeferredBases())
       .isEqualTo(ImmutableSet.of("deferred", "attribute1"));
-    assertThat(deferredToken.getSetDeferredWords())
+    assertThat(deferredToken.getSetDeferredBases())
       .isEqualTo(ImmutableSet.of("deferred", "attribute2"));
   }
 
@@ -272,9 +272,9 @@ public void itFindsFirstValidDeferredWordsWithNestedAttributes() {
       .addSetDeferredWords(ImmutableSet.of("deferred", ".attribute2.ignoreme"))
       .build();
 
-    assertThat(deferredToken.getUsedDeferredWords())
+    assertThat(deferredToken.getUsedDeferredBases())
       .isEqualTo(ImmutableSet.of("deferred", "attribute1"));
-    assertThat(deferredToken.getSetDeferredWords())
+    assertThat(deferredToken.getSetDeferredBases())
       .isEqualTo(ImmutableSet.of("deferred", "attribute2"));
   }