From 19412f007418251a915255d02d1475fa97b50ab6 Mon Sep 17 00:00:00 2001
From: Bin Wang <bin.wang@mail.binwang.me>
Date: Wed, 27 Sep 2023 12:23:53 -0400
Subject: [PATCH] Add a max size for normalize cache

This fixes #2484. Added a capacity to normalize cache. Default to 1024
and can be override by `quill.query.cacheDynamicMaxSize`.

Added a new dependency Google Guava to use its cache implementation. I
know it's better to avoid new dependency but Google Guava is a pretty
widely adopted library. It's cache implementation is widely tested. For
such a critical use case in quill, I think it's better to have a widely
tested implementation than reinventing the wheel.
---
 build.sbt                                     |  3 ++-
 .../io/getquill/norm/NormalizeCaching.scala   | 19 +++++++++----------
 .../scala/io/getquill/util/Messages.scala     |  4 ++++
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/build.sbt b/build.sbt
index 8df35756ba..15403b6e69 100644
--- a/build.sbt
+++ b/build.sbt
@@ -231,7 +231,8 @@ lazy val `quill-engine` =
         "com.typesafe.scala-logging" %% "scala-logging" % "3.9.5",
         ("com.github.takayahilton"   %% "sql-formatter" % "1.2.1").cross(CrossVersion.for3Use2_13),
         "io.suzaku"                  %% "boopickle"     % "1.4.0",
-        "com.lihaoyi"                %% "pprint"        % "0.8.1"
+        "com.lihaoyi"                %% "pprint"        % "0.8.1",
+        "com.google.guava"            % "guava"         % "32.1.2-jre"
       ),
       coverageExcludedPackages := "<empty>;.*AstPrinter;.*Using;io.getquill.Model;io.getquill.ScalarTag;io.getquill.QuotationTag"
     )
diff --git a/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala b/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala
index d7f568db6f..06dfd55b22 100644
--- a/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala
+++ b/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala
@@ -1,21 +1,20 @@
 package io.getquill.norm
 
+import com.google.common.cache.{Cache, CacheBuilder}
 import io.getquill.ast.Ast
-import java.util.concurrent.ConcurrentHashMap
+import io.getquill.util.Messages
 
 object NormalizeCaching {
-  private val cache = new ConcurrentHashMap[Ast, Ast]
+
+  private val cache: Cache[Ast, Ast] = CacheBuilder
+    .newBuilder()
+    .maximumSize(Messages.cacheDynamicMaxSize)
+    .recordStats()
+    .build()
 
   def apply(f: Ast => Ast): Ast => Ast = { ori =>
     val (stabilized, state) = StabilizeLifts.stabilize(ori)
-    val cachedR             = cache.get(stabilized)
-    val normalized = if (cachedR != null) {
-      cachedR
-    } else {
-      val r = f(stabilized)
-      cache.put(stabilized, r)
-      r
-    }
+    val normalized          = cache.get(stabilized, () => f(stabilized))
     StabilizeLifts.revert(normalized, state)
   }
 
diff --git a/quill-engine/src/main/scala/io/getquill/util/Messages.scala b/quill-engine/src/main/scala/io/getquill/util/Messages.scala
index d8adebe10c..2abc7b5b3a 100644
--- a/quill-engine/src/main/scala/io/getquill/util/Messages.scala
+++ b/quill-engine/src/main/scala/io/getquill/util/Messages.scala
@@ -49,6 +49,10 @@ object Messages {
     "quill.query.cacheDynamic",
     variable("quill.query.cacheDynamic", "query_query_cacheDynamic", "true").toBoolean
   )
+  def cacheDynamicMaxSize = cache(
+    "quill.query.cacheDynamicMaxSize",
+    variable("quill.query.cacheDynamicMaxSize", "query_query_cacheDynamicMaxSize", "1024").toLong
+  )
   def querySubexpand =
     cache("quill.query.subexpand", variable("quill.query.subexpand", "query_query_subexpand", "true").toBoolean)
   def quillLogFile = cache("quill.log.file", LogToFile(variable("quill.log.file", "quill_log_file", "false")))