From 11825c43e1ccb2fffb62e05265c839edf77173da Mon Sep 17 00:00:00 2001
From: Karol Stasiak <karol.m.stasiak@gmail.com>
Date: Wed, 17 Jul 2019 20:48:31 +0200
Subject: [PATCH] Check if all segment exist right after parsing

---
 src/main/scala/millfork/Main.scala               |  2 ++
 src/main/scala/millfork/node/Program.scala       | 12 ++++++++++++
 src/test/scala/millfork/test/emu/EmuRun.scala    |  2 ++
 src/test/scala/millfork/test/emu/EmuZ80Run.scala |  2 ++
 4 files changed, 18 insertions(+)

diff --git a/src/main/scala/millfork/Main.scala b/src/main/scala/millfork/Main.scala
index 11ca3da5..ba1803b8 100644
--- a/src/main/scala/millfork/Main.scala
+++ b/src/main/scala/millfork/Main.scala
@@ -217,6 +217,8 @@ object Main {
       OptimizationPresets.NodeOpt0.foldLeft(unoptimized)((p, opt) => p.applyNodeOptimization(opt, options))
     }
     val callGraph = new StandardCallGraph(program, options.log)
+    program.checkSegments(options.log, platform.codeAllocators.keySet)
+    options.log.assertNoErrors("Build failed due to undefined segments")
 
     val env = new Environment(None, "", platform.cpuFamily, options)
     env.collectDeclarations(program, options)
diff --git a/src/main/scala/millfork/node/Program.scala b/src/main/scala/millfork/node/Program.scala
index 20c7bd7c..c246c7d7 100644
--- a/src/main/scala/millfork/node/Program.scala
+++ b/src/main/scala/millfork/node/Program.scala
@@ -1,6 +1,7 @@
 package millfork.node
 
 import millfork.CompilationOptions
+import millfork.error.Logger
 import millfork.node.opt.NodeOptimization
 
 /**
@@ -18,6 +19,17 @@ case class Program(declarations: List[DeclarationStatement]) {
     })
   }
 
+  def checkSegments(log: Logger, existingBanks: Set[String]): Unit = {
+    declarations.foreach {
+      case s: BankedDeclarationStatement =>
+        s.bank.foreach{ b =>
+          if (!existingBanks(b)) {
+            log.error(s"Cannot allocate ${s.name}: Segment $b doesn't exist")
+          }
+        }
+      case _ =>
+    }
+  }
   def applyNodeOptimization(o: NodeOptimization, options: CompilationOptions) = Program(o.optimize(declarations, options).asInstanceOf[List[DeclarationStatement]])
   def +(p:Program): Program = Program(this.declarations ++ p.declarations)
 }
diff --git a/src/test/scala/millfork/test/emu/EmuRun.scala b/src/test/scala/millfork/test/emu/EmuRun.scala
index 6fda4b48..fd85cd53 100644
--- a/src/test/scala/millfork/test/emu/EmuRun.scala
+++ b/src/test/scala/millfork/test/emu/EmuRun.scala
@@ -184,6 +184,8 @@ class EmuRun(cpu: millfork.Cpu.Value, nodeOptimizations: List[NodeOptimization],
           tmp
         }
         val program = nodeOptimizations.foldLeft(withLibraries.applyImportantAliases)((p, opt) => p.applyNodeOptimization(opt, options))
+        program.checkSegments(log, platform.codeAllocators.keySet)
+        log.assertNoErrors("Failed")
         val callGraph = new StandardCallGraph(program, log)
         val env = new Environment(None, "", CpuFamily.M6502, options)
         env.collectDeclarations(program, options)
diff --git a/src/test/scala/millfork/test/emu/EmuZ80Run.scala b/src/test/scala/millfork/test/emu/EmuZ80Run.scala
index 5b2ff2a3..aa0a7dcc 100644
--- a/src/test/scala/millfork/test/emu/EmuZ80Run.scala
+++ b/src/test/scala/millfork/test/emu/EmuZ80Run.scala
@@ -111,6 +111,8 @@ class EmuZ80Run(cpu: millfork.Cpu.Value, nodeOptimizations: List[NodeOptimizatio
           tmp
         }
         val program = nodeOptimizations.foldLeft(withLibraries.applyImportantAliases)((p, opt) => p.applyNodeOptimization(opt, options))
+        program.checkSegments(log, platform.codeAllocators.keySet)
+        log.assertNoErrors("Failed")
         val callGraph = new StandardCallGraph(program, log)
         val env = new Environment(None, "", CpuFamily.I80, options)
         env.collectDeclarations(program, options)