From 0ae9a9f9e1cedaa54bc731df4603e0d21883ef58 Mon Sep 17 00:00:00 2001
From: hinerm <hinerm@gmail.com>
Date: Wed, 11 Dec 2024 12:20:13 -0600
Subject: [PATCH] Enable 2-stage runtime parameter processing

Previously we were performing special-case parameter processing (e.g.
expanding %'s in jvmMaxHeap specifications) at the time of initial
runtime configuration. However, this would happen a) before variable
interpolation and b) only if the user hadn't manually specified a value.
As a result, %'s would only work in a very limited number of cases (when
set in the .toml)

Instead, we now do a second "processArgs" pass that happens after
variable interpolation. It is intended to operate on all runtime args at
that point, whether they came from a toml configuration, user, etc...

This allows the particular case of % calculation in jvmMaxHeap to
trigger in all cases.

Closes #64.

Note that because we process every item in the arg list we are working
around #63, without actually fixing it.
---
 src/commonMain/kotlin/jvm.kt     | 25 +++++++++++++++++++------
 src/commonMain/kotlin/main.kt    |  9 +++++++++
 src/commonMain/kotlin/python.kt  |  4 ++++
 src/commonMain/kotlin/runtime.kt |  5 +++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/commonMain/kotlin/jvm.kt b/src/commonMain/kotlin/jvm.kt
index 8283895..0357b51 100644
--- a/src/commonMain/kotlin/jvm.kt
+++ b/src/commonMain/kotlin/jvm.kt
@@ -118,14 +118,11 @@ class JvmRuntimeConfig(recognizedArgs: Array<String>) :
             }
         }
 
-        // If not already declared, calculate and declare the max heap size.
+        // If not manually declared, add a max heap flag from the config.
         val mxIndex = runtimeArgs.indexOfFirst { it.startsWith("-Xmx") }
         if (mxIndex < 0) {
-            val maxHeap = calculateMaxHeap(config.jvmMaxHeap)
-            if (maxHeap != null) {
-                runtimeArgs += "-Xmx$maxHeap"
-                debug("Added maxHeap arg: ${runtimeArgs.last()}")
-            }
+            runtimeArgs += "-Xmx${config.jvmMaxHeap}"
+            debug("Added maxHeap arg: ${runtimeArgs.last()}")
         }
 
         // Calculate main class.
@@ -165,6 +162,22 @@ class JvmRuntimeConfig(recognizedArgs: Array<String>) :
         }
     }
 
+    override fun processArgs(args: MutableList<String>) {
+        val prefix = "-Xmx"
+        val memIndices = args.withIndex()
+            .filter{ (_, value) -> value.startsWith("-Xmx" ) }
+            .map{it.index}
+
+        for (memIdx in memIndices) {
+            if (args[memIdx].contains('%')) {
+                val memPercent = args[memIdx].substring(prefix.length)
+                val maxHeap = calculateMaxHeap(memPercent)
+                args[memIdx] = "$prefix$maxHeap";
+                debug("Expanded % in jvm runtime arg: ${args[memIdx]}")
+            }
+        }
+    }
+
     // -- Directive handlers --
 
     fun classpath(divider: String = "\n- "): String? {
diff --git a/src/commonMain/kotlin/main.kt b/src/commonMain/kotlin/main.kt
index 7dd9a8b..1c772fe 100644
--- a/src/commonMain/kotlin/main.kt
+++ b/src/commonMain/kotlin/main.kt
@@ -108,6 +108,15 @@ fun main(args: Array<String>) {
         vars.interpolateInto(programArgs.main)
     }
 
+    // Now that our program arguments have been fully interpolated, we offer the
+    // runtimes an additional opportunity to perform any runtime-specific
+    // custom logic (for example, resolving %'s in -Xmx notation).
+    for (programArgs in argsInContext.values) {
+        for (runtime in runtimes) {
+            runtime.processArgs(programArgs.runtime)
+        }
+    }
+
     // Finally, execute all the directives! \^_^/
     executeDirectives(nonGlobalDirectives, launchDirectives, runtimes, argsInContext)
 
diff --git a/src/commonMain/kotlin/python.kt b/src/commonMain/kotlin/python.kt
index e14bc53..aae3e99 100644
--- a/src/commonMain/kotlin/python.kt
+++ b/src/commonMain/kotlin/python.kt
@@ -119,6 +119,10 @@ class PythonRuntimeConfig(recognizedArgs: Array<String>) :
         }
     }
 
+    override fun processArgs(args: MutableList<String>) {
+        // No-op
+    }
+
     // -- Directive handlers --
 
     fun dryRun(args: ProgramArgs): String {
diff --git a/src/commonMain/kotlin/runtime.kt b/src/commonMain/kotlin/runtime.kt
index b238bc6..1adea28 100644
--- a/src/commonMain/kotlin/runtime.kt
+++ b/src/commonMain/kotlin/runtime.kt
@@ -43,6 +43,11 @@ abstract class RuntimeConfig(
     /** Get the launch directive block for this runtime configuration. */
     abstract fun launch(args: ProgramArgs): List<String>
 
+    /**
+     * Perform any runtime-specific argument processing here.
+     */
+    abstract fun processArgs(args: MutableList<String>)
+
     /**
      * Check whether the given argument matches one of the [recognizedArgs].
      *