Actually try to support aliased variable modification inside of macro…

… function
HubSpot · Oct 10, 2023 · ea4a7ac · ea4a7ac
1 parent fbdaf83
commit ea4a7ac
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 23 deletions.
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
@@ -1,5 +1,7 @@
 package com.hubspot.jinjava.lib.fn.eager;
 
+import static com.hubspot.jinjava.util.EagerReconstructionUtils.buildSetTag;
+
 import com.google.common.annotations.Beta;
 import com.hubspot.jinjava.el.ext.AstMacroFunction;
 import com.hubspot.jinjava.el.ext.DeferredInvocationResolutionException;
@@ -12,14 +14,18 @@
 import com.hubspot.jinjava.interpret.JinjavaInterpreter.InterpreterScopeClosable;
 import com.hubspot.jinjava.lib.fn.MacroFunction;
 import com.hubspot.jinjava.lib.tag.MacroTag;
+import com.hubspot.jinjava.lib.tag.eager.DeferredToken;
 import com.hubspot.jinjava.lib.tag.eager.EagerExecutionResult;
+import com.hubspot.jinjava.lib.tag.eager.importing.AliasedEagerImportingStrategy;
 import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
 import com.hubspot.jinjava.tree.Node;
+import com.hubspot.jinjava.tree.parse.TagToken;
 import com.hubspot.jinjava.util.EagerContextWatcher;
 import com.hubspot.jinjava.util.EagerContextWatcher.EagerChildContextConfig;
 import com.hubspot.jinjava.util.EagerExpressionResolver.EagerExpressionResult;
 import com.hubspot.jinjava.util.EagerReconstructionUtils;
 import com.hubspot.jinjava.util.PrefixToPreserveState;
+import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,6 +36,7 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.apache.commons.lang3.StringUtils;
 
 @Beta
@@ -71,15 +78,15 @@ public Object doEvaluate(
           interpreter
         );
         if (!result.getResult().isFullyResolved()) {
-          if (
-            result
-              .getSpeculativeBindings()
-              .keySet()
-              .stream()
-              .anyMatch(key -> localContextScope.getScope().containsKey(key))
-          ) {
-            throw new DeferredValueException("e");
-          }
+          //          if (
+          //            result
+          //              .getSpeculativeBindings()
+          //              .keySet()
+          //              .stream()
+          //              .anyMatch(key -> localContextScope.getScope().containsKey(key))
+          //          ) {
+          //            throw new DeferredValueException("e");
+          //          }
           result =
             eagerEvaluateInDeferredExecutionMode(
               () -> getEvaluationResultDirectly(argMap, kwargMap, varArgs, interpreter),
@@ -116,6 +123,47 @@ public Object doEvaluate(
         interpreter,
         eagerExecutionResult
       );
+      Optional<String> maybeTempImportAlias = AliasedEagerImportingStrategy.getTemporaryImportAlias(
+        localContextScope
+      );
+      Optional<String> maybeActualImportAlias = localContextScope.getImportResourceAlias();
+      if (
+        maybeTempImportAlias.isPresent() &&
+        maybeActualImportAlias.isPresent() &&
+        interpreter
+          .getContext()
+          .getDeferredTokens()
+          .stream()
+          .anyMatch(
+            deferredToken ->
+              deferredToken.getUsedDeferredWords().contains(maybeTempImportAlias.get())
+          )
+      ) {
+        String reconstructTemporaryAliasSetTag = buildSetTag(
+          Collections.singletonMap(
+            maybeTempImportAlias.get(),
+            maybeActualImportAlias.get()
+          ),
+          interpreter,
+          false
+        );
+        EagerReconstructionUtils.handleDeferredTokenAndReconstructReferences(
+          interpreter,
+          DeferredToken
+            .builderFromImage(
+              reconstructTemporaryAliasSetTag,
+              TagToken.class,
+              interpreter
+            )
+            .addUsedDeferredWords(Stream.of(maybeActualImportAlias.get()))
+            .addSetDeferredWords(Stream.of(maybeTempImportAlias.get()))
+            .build()
+        );
+        prefixToPreserveState.put(
+          maybeTempImportAlias.get(),
+          reconstructTemporaryAliasSetTag
+        );
+      }
 
       String tempVarName = MacroFunctionTempVariable.getVarName(
         getName(),
@@ -303,7 +351,7 @@ private String getSetTagForAliasedVariables(String fullName) {
         .getCombinedScope()
         .entrySet()
         .stream()
-        //        .filter(entry -> entry.getValue() instanceof DeferredValue)
+        .filter(entry -> entry.getValue() instanceof DeferredValue)
         .map(Entry::getKey)
         .collect(Collectors.toMap(Function.identity(), name -> aliasName + name));
       return EagerReconstructionUtils.buildSetTag(

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerSetTagStrategy.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/EagerSetTagStrategy.java
@@ -5,6 +5,7 @@
 import com.hubspot.jinjava.lib.tag.SetTag;
 import com.hubspot.jinjava.lib.tag.eager.importing.AliasedEagerImportingStrategy;
 import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.tree.parse.TagToken;
 import com.hubspot.jinjava.util.EagerReconstructionUtils;
 import com.hubspot.jinjava.util.PrefixToPreserveState;
 import java.util.Arrays;
@@ -166,16 +167,24 @@ public static String getSuffixToPreserveState(
       !AliasedEagerImportingStrategy.isTemporaryImportAlias(variables) &&
       !interpreter.getContext().getMetaContextVariables().contains(variables)
     ) {
-      String updateString = getUpdateString(variables);
-
+      if (!interpreter.getContext().containsKey(maybeTemporaryImportAlias.get())) {
+        maybeTemporaryImportAlias = interpreter.getContext().getImportResourceAlias();
+      }
       // Don't need to render because the temporary import alias's value is always deferred, and rendering will do nothing
-      suffixToPreserveState.append(
-        EagerReconstructionUtils.buildDoUpdateTag(
-          maybeTemporaryImportAlias.get(),
-          updateString,
-          interpreter
-        )
+
+      String doUpdateImage = EagerReconstructionUtils.buildDoUpdateTag(
+        maybeTemporaryImportAlias.get(),
+        getUpdateString(variables),
+        interpreter
+      );
+      EagerReconstructionUtils.handleDeferredTokenAndReconstructReferences(
+        interpreter,
+        DeferredToken
+          .builderFromImage(doUpdateImage, TagToken.class, interpreter)
+          .addUsedDeferredWords(Stream.of(maybeTemporaryImportAlias.get()))
+          .build()
       );
+      suffixToPreserveState.append(doUpdateImage);
     }
 
     return suffixToPreserveState.toString();

diff --git a/src/main/java/com/hubspot/jinjava/lib/tag/eager/importing/AliasedEagerImportingStrategy.java b/src/main/java/com/hubspot/jinjava/lib/tag/eager/importing/AliasedEagerImportingStrategy.java
@@ -112,12 +112,24 @@ public void integrateChild(JinjavaInterpreter child) {
       currentImportAlias,
       child
     );
-    childBindings.remove(temporaryImportAlias);
-    importingData.getOriginalInterpreter().getContext().remove(temporaryImportAlias);
+    //    childBindings.remove(temporaryImportAlias);
+    //    importingData.getOriginalInterpreter().getContext().remove(temporaryImportAlias);
     // Remove meta keys
-    childBindings.remove(Context.GLOBAL_MACROS_SCOPE_KEY);
-    childBindings.remove(Context.IMPORT_RESOURCE_ALIAS_KEY);
-    mapForCurrentContextAlias.putAll(childBindings);
+    //    childBindings.remove(Context.GLOBAL_MACROS_SCOPE_KEY);
+    //    childBindings.remove(Context.IMPORT_RESOURCE_ALIAS_KEY);
+    childBindings
+      .entrySet()
+      .stream()
+      .filter(
+        entry ->
+          !(
+            entry.getKey().equals(Context.GLOBAL_MACROS_SCOPE_KEY) ||
+            entry.getKey().equals(Context.IMPORT_RESOURCE_ALIAS_KEY) ||
+            AliasedEagerImportingStrategy.isTemporaryImportAlias(entry.getKey())
+          )
+      )
+      .forEach(entry -> mapForCurrentContextAlias.put(entry.getKey(), entry.getValue()));
+    //    mapForCurrentContextAlias.putAll(childBindings);
   }
 
   @Override