feat: Add optimizable function (#204)

core/src/fr/hammons/slinc/jitc/OptimizableFunction.scala

@@ -0,0 +1,36 @@
+package fr.hammons.slinc.jitc
+
+import java.util.concurrent.atomic.AtomicInteger
+import scala.compiletime.codeOf
+
+sealed trait OptimizableFunction[F]:
+  def apply[O](fn: F => O): O
+
+class UnoptimizedFunction[F](
+    aotcFn: F,
+    val originalCode: String,
+    limit: Int,
+    updateFn: () => Unit
+) extends OptimizableFunction[F]:
+  val counter: AtomicInteger = AtomicInteger(0)
+
+  def getCount() = counter.getOpaque()
+
+  def apply[O](fn: F => O): O =
+    val res = fn(aotcFn)
+    val currentCount = counter.getOpaque()
+    if currentCount + 1 >= limit then
+      counter.setOpaque(currentCount + 1)
+      updateFn()
+      res
+    else
+      counter.setOpaque(currentCount + 1)
+      res
+
+object UnoptimizedFunction:
+  inline def apply[F](
+      inline aotc: F,
+      limit: Int,
+      updateFn: () => Unit
+  ): UnoptimizedFunction[F] =
+    new UnoptimizedFunction[F](aotc, codeOf(aotc), limit, updateFn)

core/test/src/fr/hammons/slinc/jitc/JitSpecification.scala

@@ -0,0 +1,58 @@
+package fr.hammons.slinc.jitc
+
+import scala.concurrent.Future
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
+import scala.concurrent.ExecutionContext.Implicits.global
+import munit.ScalaCheckSuite
+import org.scalacheck.Prop.*
+import org.scalacheck.Gen
+
+class JitSpecification extends ScalaCheckSuite:
+  test("single-threaded count works"):
+      var optimizationTriggered = false
+      val fn = UnoptimizedFunction(
+        (i: Int) => i,
+        10,
+        () => optimizationTriggered = true
+      )
+
+      var i = 0
+      while i < 10 do
+        fn(_(6))
+        i += 1
+
+      assertEquals(fn.getCount(), 10)
+      assert(
+        optimizationTriggered,
+        s"Optimization function not run? ${optimizationTriggered} - ${fn.getCount()}"
+      )
+      assertNoDiff(
+        fn.originalCode,
+        """|{
+       |  def $anonfun(i: Int): Int = i
+       |  closure($anonfun)
+       |}
+       |""".stripMargin
+      )
+
+  property("multi-threaded count works"):
+      forAll(Gen.choose(0, 100)): (runs: Int) =>
+        var optimizationTriggered = false
+        val fn = UnoptimizedFunction(
+          (i: Int) => i,
+          10,
+          () => optimizationTriggered = true
+        )
+
+        val futures: Seq[Future[Unit]] = for
+          j <- 0 until 5
+          res = Future(
+            (0 until runs).foreach(i => fn(_(i * j)))
+          )
+        yield res
+
+        Await.result(Future.traverse(futures)(identity), Duration.Inf)
+
+        assert(fn.getCount() >= runs)
+        assert(runs < 10 || optimizationTriggered)