From dc93691fd9993a7340d2090d42c443e062ed4895 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Sun, 6 Mar 2022 20:35:15 +0100
Subject: [PATCH] working on new ast

---
 .../src/prog8/codegen/experimental6502/AsmGen.kt           | 3 +++
 compiler/test/ast/TestIntermediateAst.kt                   | 7 +++++++
 .../src/prog8/compilerinterface/intermediate/AstBase.kt    | 6 ++++++
 3 files changed, 16 insertions(+)

diff --git a/codeGenExperimental6502/src/prog8/codegen/experimental6502/AsmGen.kt b/codeGenExperimental6502/src/prog8/codegen/experimental6502/AsmGen.kt
index 0707055f8..ef9f83100 100644
--- a/codeGenExperimental6502/src/prog8/codegen/experimental6502/AsmGen.kt
+++ b/codeGenExperimental6502/src/prog8/codegen/experimental6502/AsmGen.kt
@@ -1,6 +1,7 @@
 package prog8.codegen.experimental6502
 
 import prog8.ast.Program
+import prog8.ast.base.FatalAstException
 import prog8.compilerinterface.*
 
 class AsmGen(internal val program: Program,
@@ -17,6 +18,8 @@ class AsmGen(internal val program: Program,
         // TODO temporary location to do this:
         val intermediateAst = IntermediateAstMaker.transform(program)
         intermediateAst.print()
+        val entry = intermediateAst.entrypoint() ?: throw FatalAstException("no main.start() found")
+        println(entry)
 
         println("..todo: create assembly code into ${options.outputDir.toAbsolutePath()}..")
         return AssemblyProgram("dummy")
diff --git a/compiler/test/ast/TestIntermediateAst.kt b/compiler/test/ast/TestIntermediateAst.kt
index 67ed7e4f7..dc1bd2c35 100644
--- a/compiler/test/ast/TestIntermediateAst.kt
+++ b/compiler/test/ast/TestIntermediateAst.kt
@@ -1,6 +1,8 @@
 package prog8tests.ast
 
+import io.kotest.assertions.fail
 import io.kotest.core.spec.style.FunSpec
+import io.kotest.matchers.ints.shouldBeGreaterThan
 import io.kotest.matchers.shouldBe
 import prog8.codegen.experimental6502.IntermediateAstMaker
 import prog8.codegen.target.C64Target
@@ -26,6 +28,11 @@ class TestIntermediateAst: FunSpec({
         val ast = IntermediateAstMaker.transform(result.program)
         ast.name shouldBe result.program.name
         ast.builtinFunctions.names shouldBe result.program.builtinFunctions.names
+        val entry = ast.entrypoint() ?: fail("no main.start() found")
+        entry.name shouldBe "start"
+        val blocks = ast.allBlocks().toList()
+        blocks.size shouldBeGreaterThan 1
+        blocks[0].name shouldBe "main"
         ast.print()
     }
 
diff --git a/compilerInterfaces/src/prog8/compilerinterface/intermediate/AstBase.kt b/compilerInterfaces/src/prog8/compilerinterface/intermediate/AstBase.kt
index 3b2679c97..c36880663 100644
--- a/compilerInterfaces/src/prog8/compilerinterface/intermediate/AstBase.kt
+++ b/compilerInterfaces/src/prog8/compilerinterface/intermediate/AstBase.kt
@@ -50,6 +50,12 @@ class PtProgram(
     override fun printProperties() {
         print("'$name'")
     }
+
+    fun allBlocks(): Sequence<PtBlock> =
+        children.asSequence().flatMap { it.children }.filterIsInstance<PtBlock>()
+
+    fun entrypoint(): PtSub? =
+        allBlocks().firstOrNull { it.name == "main" }?.children?.firstOrNull { it is PtSub && it.name == "start" } as PtSub?
 }