From f5aba87f506c0bb7619d76168b52ae9eacc3af37 Mon Sep 17 00:00:00 2001
From: Edwin Jakobs <edwin@rndr.studio>
Date: Sun, 3 May 2020 12:54:32 +0200
Subject: [PATCH] Improve support for olive programs inside packages

---
 .../src/main/kotlin/ExtractProgram.kt            | 16 ++++++++++++++--
 orx-olive/src/main/kotlin/OliveProgram.kt        |  2 +-
 orx-olive/src/main/kotlin/ScriptGenerator.kt     |  2 ++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/orx-kotlin-parser/src/main/kotlin/ExtractProgram.kt b/orx-kotlin-parser/src/main/kotlin/ExtractProgram.kt
index 91319f59..c2f0a8de 100644
--- a/orx-kotlin-parser/src/main/kotlin/ExtractProgram.kt
+++ b/orx-kotlin-parser/src/main/kotlin/ExtractProgram.kt
@@ -18,6 +18,13 @@ fun ParserRuleContext.verbatimText(marginLeft: Int = 0, marginRight: Int = 0): S
     return start.inputStream.getText(interval)
 }
 
+class PackageExtractor() : KotlinParserBaseListener() {
+    var result: String? = null
+    override fun enterPackageHeader(ctx: KotlinParser.PackageHeaderContext) {
+        result = ctx.verbatimText()
+    }
+}
+
 class ImportsExtractor(val ruleNames: List<String>) : KotlinParserBaseListener() {
     var result: String? = null
 
@@ -30,6 +37,7 @@ class LambdaExtractor(val ruleNames: List<String>, val lambdaName: String) : Kot
     fun RuleContext.named(): String {
         return ruleNames[this.ruleIndex]
     }
+
     var result: String? = null
     override fun enterAnnotatedLambda(ctx: KotlinParser.AnnotatedLambdaContext?) {
         val puec = ctx!!.parent!!.parent!!.parent!! as KotlinParser.PostfixUnaryExpressionContext
@@ -42,7 +50,7 @@ class LambdaExtractor(val ruleNames: List<String>, val lambdaName: String) : Kot
     }
 }
 
-class ProgramSource(val imports: String, val programLambda: String)
+class ProgramSource(val packageName: String?, val imports: String, val programLambda: String)
 
 fun extractProgram(source: String, programIdentifier: String = "program"): ProgramSource {
     val parser = KotlinParser(CommonTokenStream(KotlinLexer(CharStreams.fromString(source))))
@@ -50,10 +58,14 @@ fun extractProgram(source: String, programIdentifier: String = "program"): Progr
 //    val rules = parser.ruleNames.toList()
 //    val pt = TreeUtils.toPrettyTree(root, rules)
     val ruleNames = parser.ruleNames.toList()
+
+    val packageExtractor = PackageExtractor()
+    ParseTreeWalker.DEFAULT.walk(packageExtractor, root)
+
     val importsExtractor = ImportsExtractor(ruleNames)
     ParseTreeWalker.DEFAULT.walk(importsExtractor, root)
 
     val lambdaExtractor = LambdaExtractor(ruleNames, programIdentifier)
     ParseTreeWalker.DEFAULT.walk(lambdaExtractor, root)
-    return ProgramSource(importsExtractor.result!!, lambdaExtractor.result!!)
+    return ProgramSource(packageExtractor.result, importsExtractor.result?:"", lambdaExtractor.result?:"")
 }
diff --git a/orx-olive/src/main/kotlin/OliveProgram.kt b/orx-olive/src/main/kotlin/OliveProgram.kt
index e74e9227..45b02f88 100644
--- a/orx-olive/src/main/kotlin/OliveProgram.kt
+++ b/orx-olive/src/main/kotlin/OliveProgram.kt
@@ -20,7 +20,7 @@ fun stackRootClassName(thread: Thread = Thread.currentThread(), sanitize: Boolea
 }
 
 fun ApplicationBuilder.oliveProgram(init: OliveProgram.() -> Unit): OliveProgram {
-    val rootClassName = stackRootClassName(sanitize = true)
+    val rootClassName = stackRootClassName(sanitize = true).split(".").last()
 
     var sourceLocation = "src/main/kotlin/$rootClassName.kt"
     val candidateFile = File(sourceLocation)
diff --git a/orx-olive/src/main/kotlin/ScriptGenerator.kt b/orx-olive/src/main/kotlin/ScriptGenerator.kt
index a62f0c82..17b41dfc 100644
--- a/orx-olive/src/main/kotlin/ScriptGenerator.kt
+++ b/orx-olive/src/main/kotlin/ScriptGenerator.kt
@@ -4,7 +4,9 @@ import org.openrndr.extra.kotlinparser.ProgramSource
 
 inline fun <reified T> generateScript(programSource: ProgramSource): String {
     val script = """
+        
 @file:Suppress("UNUSED_LAMBDA_EXPRESSION")
+${programSource.packageName?:""}
 
 import org.openrndr.extra.olive.OliveProgram
 ${programSource.imports}