From 3a56a7114dc010fbd7a6bbd572095fe9e5f6ff2c Mon Sep 17 00:00:00 2001
From: davidmarne-wf <david.marne@workiva.com>
Date: Thu, 16 Dec 2021 22:45:26 -0800
Subject: [PATCH] compile: fix file renaming check in optimization (#4154)

This was resulting in a max of 2 optimized entrypoints in a single file.

Signed-off-by: davidmarne-wf <david.marne@workiva.com>
---
 compile/compile.go      |  7 ++++---
 compile/compile_test.go | 11 ++++++++++-
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/compile/compile.go b/compile/compile.go
index 70177f2499..9d7fdca94d 100644
--- a/compile/compile.go
+++ b/compile/compile.go
@@ -862,12 +862,13 @@ func (o *optimizer) getSupportModuleFilename(used map[string]int, module *ast.Mo
 
 	if err == nil && safePathPattern.MatchString(fileName) {
 		fileName = o.outputprefix + "/" + fileName
+		uniqueFileName := fileName
 		if c, ok := used[fileName]; ok {
-			fileName += fmt.Sprintf(".%d", c)
+			uniqueFileName += fmt.Sprintf(".%d", c)
 		}
 		used[fileName]++
-		fileName += ".rego"
-		return fileName
+		uniqueFileName += ".rego"
+		return uniqueFileName
 	}
 
 	return fmt.Sprintf("%v/%v/%v/%v.rego", o.outputprefix, o.nsprefix, entrypointIndex, supportIndex)
diff --git a/compile/compile_test.go b/compile/compile_test.go
index 0cd1ee0fd4..4b6524dfbe 100644
--- a/compile/compile_test.go
+++ b/compile/compile_test.go
@@ -891,7 +891,7 @@ func TestOptimizerOutput(t *testing.T) {
 		},
 		{
 			note:        "multiple entrypoints",
-			entrypoints: []string{"data.test.p", "data.test.r"},
+			entrypoints: []string{"data.test.p", "data.test.r", "data.test.s"},
 			modules: map[string]string{
 				"test.rego": `
 					package test
@@ -904,6 +904,10 @@ func TestOptimizerOutput(t *testing.T) {
 						q[input.x]
 					}
 
+					s {
+						q[input.x]
+					}
+
 					q[1]
 				`,
 			},
@@ -918,6 +922,11 @@ func TestOptimizerOutput(t *testing.T) {
 
 					r = __result__ { 1 = input.x; __result__ = true }
 				`,
+				"optimized/test.2.rego": `
+					package test
+
+					s = __result__ { 1 = input.x; __result__ = true }
+				`,
 				"test.rego": `
 					package test