From 66fde129c7f4fef3b70a1ab8dcfff14c0d49abc9 Mon Sep 17 00:00:00 2001
From: Chinmay Madeshi <madeshic@google.com>
Date: Thu, 18 Sep 2025 10:29:59 -0700
Subject: [PATCH] Internal Changes

PiperOrigin-RevId: 808640842
---
 .../testing/testrunner/CelCoverageIndex.java  | 22 ++++++++++-------
 .../testrunner/CelCoverageIndexTest.java      | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/testing/src/main/java/dev/cel/testing/testrunner/CelCoverageIndex.java b/testing/src/main/java/dev/cel/testing/testrunner/CelCoverageIndex.java
index e7b08da79..72b146516 100644
--- a/testing/src/main/java/dev/cel/testing/testrunner/CelCoverageIndex.java
+++ b/testing/src/main/java/dev/cel/testing/testrunner/CelCoverageIndex.java
@@ -61,15 +61,19 @@ final class CelCoverageIndex {
       new ConcurrentHashMap<>();
 
   public void init(CelAbstractSyntaxTree ast) {
-    this.ast = ast;
-    CelNavigableExpr.fromExpr(ast.getExpr())
-        .allNodes()
-        .forEach(
-            celNavigableExpr -> {
-              NodeCoverageStats nodeCoverageStats = new NodeCoverageStats();
-              nodeCoverageStats.isBooleanNode.set(isNodeTypeBoolean(celNavigableExpr.expr()));
-              nodeCoverageStatsMap.put(celNavigableExpr.id(), nodeCoverageStats);
-            });
+    // If the AST and node coverage stats map are already initialized, then we don't need to
+    // re-initialize them.
+    if (this.ast == null && nodeCoverageStatsMap.isEmpty()) {
+      this.ast = ast;
+      CelNavigableExpr.fromExpr(ast.getExpr())
+          .allNodes()
+          .forEach(
+              celNavigableExpr -> {
+                NodeCoverageStats nodeCoverageStats = new NodeCoverageStats();
+                nodeCoverageStats.isBooleanNode.set(isNodeTypeBoolean(celNavigableExpr.expr()));
+                nodeCoverageStatsMap.put(celNavigableExpr.id(), nodeCoverageStats);
+              });
+    }
   }
 
   /**
diff --git a/testing/src/test/java/dev/cel/testing/testrunner/CelCoverageIndexTest.java b/testing/src/test/java/dev/cel/testing/testrunner/CelCoverageIndexTest.java
index 287d5435e..4c6a6cf26 100644
--- a/testing/src/test/java/dev/cel/testing/testrunner/CelCoverageIndexTest.java
+++ b/testing/src/test/java/dev/cel/testing/testrunner/CelCoverageIndexTest.java
@@ -154,4 +154,28 @@ public void getCoverageReport_fullCoverage_writesToUndeclaredOutputs() throws Ex
     String fileContent = Files.asCharSource(outputFile, UTF_8).read();
     assertThat(fileContent).isEqualTo(report.dotGraph());
   }
+
+  @Test
+  public void getCoverageReport_fullCoverage_multipleEvaluations() throws Exception {
+    Cel cel = CelFactory.standardCelBuilder().addVar("x", SimpleType.INT).build();
+    CelAbstractSyntaxTree ast = cel.compile("x > 1").getAst();
+    CelRuntime.Program program = cel.createProgram(ast);
+    CelCoverageIndex coverageIndex = new CelCoverageIndex();
+    coverageIndex.init(ast);
+    CelEvaluationListener listener = coverageIndex.newEvaluationListener();
+
+    program.trace(ImmutableMap.of("x", 2L), listener);
+    coverageIndex.init(ast); // Re-initialize the coverage index.
+    program.trace(ImmutableMap.of("x", 0L), listener);
+
+    CoverageReport report = coverageIndex.generateCoverageReport();
+    assertThat(report.nodes()).isGreaterThan(0);
+    assertThat(report.coveredNodes()).isEqualTo(report.nodes());
+    assertThat(report.branches()).isEqualTo(2);
+    // Despite re-initializing the coverage index now, the report should still
+    // be fully covered. Else, only the second evaluation would've been covered.
+    assertThat(report.coveredBooleanOutcomes()).isEqualTo(2);
+    assertThat(report.unencounteredNodes()).isEmpty();
+    assertThat(report.unencounteredBranches()).isEmpty();
+  }
 }