diff --git a/extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java b/extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java
index 8b67d5c79..85ee9c756 100644
--- a/extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java
+++ b/extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java
@@ -297,6 +297,7 @@ static CelExtensionLibrary<CelOptionalLibrary> library() {
   public static final CelOptionalLibrary INSTANCE = CelOptionalLibrary.library().latest();
 
   private static final String UNUSED_ITER_VAR = "#unused";
+  private static final String OPTIONAL_MAP_VAR = "@target";
 
   private final int version;
   private final ImmutableSet<CelFunctionDecl> functions;
@@ -524,21 +525,51 @@ private static Optional<CelExpr> expandOptMap(
     CelExpr mapExpr = checkNotNull(arguments.get(1));
     String varName = varIdent.ident().name();
 
-    return Optional.of(
+    if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
+      return Optional.of(
+          exprFactory.newGlobalCall(
+              Operator.CONDITIONAL.getFunction(),
+              exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
+              exprFactory.newGlobalCall(
+                  OPTIONAL_OF.getFunction(),
+                  exprFactory.fold(
+                      UNUSED_ITER_VAR,
+                      exprFactory.newList(),
+                      varName,
+                      exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
+                      exprFactory.newBoolLiteral(true),
+                      exprFactory.newIdentifier(varName),
+                      mapExpr)),
+              exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
+    }
+
+    CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
+    CelExpr localVarCopy = exprFactory.copy(localVar);
+    CelExpr conditionalExpr =
         exprFactory.newGlobalCall(
             Operator.CONDITIONAL.getFunction(),
-            exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
+            exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
             exprFactory.newGlobalCall(
                 OPTIONAL_OF.getFunction(),
                 exprFactory.fold(
                     UNUSED_ITER_VAR,
                     exprFactory.newList(),
                     varName,
-                    exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
+                    exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
                     exprFactory.newBoolLiteral(true),
                     exprFactory.newIdentifier(varName),
                     mapExpr)),
-            exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
+            exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
+
+    return Optional.of(
+        exprFactory.fold(
+            UNUSED_ITER_VAR,
+            exprFactory.newList(),
+            OPTIONAL_MAP_VAR,
+            target,
+            exprFactory.newBoolLiteral(false),
+            exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
+            conditionalExpr));
   }
 
   private static Optional<CelExpr> expandOptFlatMap(
@@ -558,19 +589,47 @@ private static Optional<CelExpr> expandOptFlatMap(
     CelExpr mapExpr = checkNotNull(arguments.get(1));
     String varName = varIdent.ident().name();
 
-    return Optional.of(
+    if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
+      return Optional.of(
+          exprFactory.newGlobalCall(
+              Operator.CONDITIONAL.getFunction(),
+              exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
+              exprFactory.fold(
+                  UNUSED_ITER_VAR,
+                  exprFactory.newList(),
+                  varName,
+                  exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
+                  exprFactory.newBoolLiteral(true),
+                  exprFactory.newIdentifier(varName),
+                  mapExpr),
+              exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
+    }
+
+    CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
+    CelExpr localVarCopy = exprFactory.copy(localVar);
+    CelExpr conditionalExpr =
         exprFactory.newGlobalCall(
             Operator.CONDITIONAL.getFunction(),
-            exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
+            exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
             exprFactory.fold(
                 UNUSED_ITER_VAR,
                 exprFactory.newList(),
                 varName,
-                exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
+                exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
                 exprFactory.newBoolLiteral(true),
                 exprFactory.newIdentifier(varName),
                 mapExpr),
-            exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
+            exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
+
+    return Optional.of(
+        exprFactory.fold(
+            UNUSED_ITER_VAR,
+            exprFactory.newList(),
+            OPTIONAL_MAP_VAR,
+            target,
+            exprFactory.newBoolLiteral(false),
+            exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
+            conditionalExpr));
   }
 
   private static Object indexOptionalMap(
diff --git a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
index 9fda186cf..920ba537b 100644
--- a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
+++ b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
@@ -17,6 +17,7 @@ java_library(
         "//common:compiler_common",
         "//common:container",
         "//common:options",
+        "//common/ast",
         "//common/exceptions:attribute_not_found",
         "//common/exceptions:divide_by_zero",
         "//common/exceptions:index_out_of_bounds",
diff --git a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
index 2ba12910f..fab444750 100644
--- a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
+++ b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
@@ -34,6 +34,7 @@
 import dev.cel.common.CelOverloadDecl;
 import dev.cel.common.CelValidationException;
 import dev.cel.common.CelVarDecl;
+import dev.cel.common.ast.CelExpr;
 import dev.cel.common.types.CelType;
 import dev.cel.common.types.ListType;
 import dev.cel.common.types.MapType;
@@ -1571,6 +1572,68 @@ public void optionalFlatMapMacro_receiverHasValue_returnsOptionalValue() throws
     assertThat(result).hasValue(43L);
   }
 
+  @Test
+  public void optionalMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .addVar("x", OptionalType.create(SimpleType.INT))
+            .setResultType(OptionalType.create(SimpleType.INT))
+            .build();
+    CelAbstractSyntaxTree ast = compile(cel, "x.optMap(y, y + 1)");
+
+    assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
+  }
+
+  @Test
+  public void optionalMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .setResultType(OptionalType.create(SimpleType.INT))
+            .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
+            .build();
+    CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optMap(y, y + 1)");
+
+    assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
+    assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
+
+    Optional<Long> result =
+        (Optional<Long>)
+            cel.createProgram(ast)
+                .eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
+    assertThat(result).hasValue(43L);
+  }
+
+  @Test
+  public void optionalFlatMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .addVar("x", OptionalType.create(SimpleType.INT))
+            .setResultType(OptionalType.create(SimpleType.INT))
+            .build();
+    CelAbstractSyntaxTree ast = compile(cel, "x.optFlatMap(y, optional.of(y + 1))");
+
+    assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
+  }
+
+  @Test
+  public void optionalFlatMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .setResultType(OptionalType.create(SimpleType.INT))
+            .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
+            .build();
+    CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optFlatMap(y, optional.of(y + 1))");
+
+    assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
+    assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
+
+    Optional<Long> result =
+        (Optional<Long>)
+            cel.createProgram(ast)
+                .eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
+    assertThat(result).hasValue(43L);
+  }
+
   @Test
   public void optionalFlatMapMacro_withOptionalOfNonZeroValue_optionalEmptyWhenValueIsZero()
       throws Exception {